Skip to main content

orch/
control.rs

1//! Control session — context assembly, one-shot agent invocation, response parsing.
2//!
3//! The control session provides an interactive ops assistant that assembles context
4//! from SQLite (memories, summaries, live state), invokes an agent CLI one-shot
5//! with `--session-id` for conversation continuity, and parses the response for
6//! summary extraction.
7//!
8//! ## Session continuity
9//!
10//! Each named session (e.g., "default", "ops") gets a UUID stored in SQLite KV
11//! (`control:session_id:{name}`). Session agents (claude, kimi, minimax, grok)
12//! receive `--session-id <uuid>` so they resume prior conversation state.
13//! The UUID is regenerated when the agent or model changes.
14//!
15//! ## Model selection
16//!
17//! `/model <spec>` where spec is:
18//! - `agent:model` — explicit agent + model (e.g., `minimax:sonnet`, `opencode:minimax-m2.5-free`)
19//! - `model` — infer agent from model name (e.g., `sonnet` → claude, `gpt-4o` → codex)
20//!
21//! All agents: claude, codex, opencode, kimi, minimax, grok
22//! - kimi and minimax are claude-compatible wrappers (separate binaries in PATH)
23//! - opencode supports `opencode models` for listing available models
24
25use std::collections::HashMap;
26use std::sync::{Arc, LazyLock, Mutex};
27
28use anyhow::{Context, Result};
29use regex::Regex;
30use tokio::time::{timeout, Duration};
31use uuid::Uuid;
32
33const SUBPROCESS_TIMEOUT: Duration = Duration::from_secs(10);
34
35use crate::store::TaskStore;
36
37/// Per-session invocation locks — ensures only one agent invocation runs at a time per session.
38///
39/// Keyed by `session_id`. A new `Arc<tokio::sync::Mutex<()>>` is inserted on first use and
40/// reused for all subsequent calls with the same session_id. Entries are evicted after each
41/// `send_message` call when no concurrent waiters remain. Note: because callers keep a
42/// local `Arc` clone while executing, the map entry + local clone equal 2 strong refs when
43/// there are no other waiters. We therefore test for `Arc::strong_count() == 2` when
44/// deciding to evict the map entry (map entry + local clone). If other waiters exist the
45/// count will be > 2 and we leave the entry in place.
46static SESSION_LOCKS: std::sync::LazyLock<Mutex<HashMap<String, Arc<tokio::sync::Mutex<()>>>>> =
47    std::sync::LazyLock::new(|| Mutex::new(HashMap::new()));
48
49/// System prompt template, loaded at compile time.
50const SYSTEM_TEMPLATE: &str = include_str!("../prompts/control_system.md");
51
52/// Timeout for agent invocations (2 minutes).
53const AGENT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(120);
54
55/// Maximum bytes for a single control-session turn payload.
56/// When the combined state + user message exceeds this limit, the payload is split
57/// across multiple sequential turns so the user's question is never truncated.
58const MAX_CONTROL_MESSAGE_BYTES: usize = 32 * 1024;
59
60/// Delivery mode for control session messages.
61enum DeliveryMode {
62    /// Single message delivery (under-limit messages)
63    Single,
64    /// Chunked delivery with session continuity (claude, kimi, minimax) — uses --resume
65    Chunked,
66    /// Chunked delivery without session continuity (codex, opencode) — independent invocations
67    ChunkedNonSession,
68}
69
70/// Parsed model specification from `/model` command.
71#[derive(Debug, Clone, PartialEq)]
72pub struct ModelSpec {
73    pub agent: String,
74    pub model: String,
75}
76
77/// Parsed control command from the chat session.
78#[derive(Debug, Clone, PartialEq, Eq)]
79pub enum ControlCommand {
80    Model(Option<String>),
81    Agent(Option<String>),
82    Restart,
83}
84
85fn default_model_for_agent(agent: &str) -> &'static str {
86    match agent {
87        "codex" => "gpt-4o",
88        "opencode" => "deepseek-r1",
89        // Empty means omit --model and let the grok CLI use its configured default
90        "grok" => "",
91        _ => "sonnet",
92    }
93}
94
95/// Parse a control-session command.
96pub fn parse_control_command(message: &str) -> Option<ControlCommand> {
97    let trimmed = message.trim();
98    let rest = trimmed.strip_prefix('/')?;
99    let (command, args) = rest
100        .split_once(' ')
101        .map_or((rest, ""), |(cmd, tail)| (cmd, tail.trim()));
102
103    match command {
104        "model" => Some(ControlCommand::Model(
105            (!args.is_empty()).then(|| args.to_string()),
106        )),
107        "agent" => Some(ControlCommand::Agent(
108            (!args.is_empty()).then(|| args.to_string()),
109        )),
110        "restart" => Some(ControlCommand::Restart),
111        _ => None,
112    }
113}
114
115/// Parse a model spec string into agent + model.
116///
117/// Formats:
118/// - `agent:model` → explicit (e.g., `minimax:sonnet`, `opencode:minimax-m2.5-free`)
119/// - `model` → infer agent from model name
120pub fn parse_model_spec(spec: &str) -> ModelSpec {
121    if let Some((agent, model)) = spec.split_once(':') {
122        ModelSpec {
123            agent: agent.to_string(),
124            model: crate::engine::router::RouterConfig::normalize_model_identifier(model),
125        }
126    } else {
127        ModelSpec {
128            agent: infer_agent(spec).to_string(),
129            model: crate::engine::router::RouterConfig::normalize_model_identifier(spec),
130        }
131    }
132}
133
134/// Infer which agent CLI should run a given model name.
135fn infer_agent(model: &str) -> &'static str {
136    let lower = model.to_lowercase();
137    if lower.contains("gpt")
138        || lower.contains("o1")
139        || lower.contains("o3")
140        || lower.contains("o4")
141        || lower.contains("codex")
142    {
143        "codex"
144    } else if lower.contains("deepseek") || lower.contains("qwen") {
145        "opencode"
146    } else if lower.contains("grok") {
147        "grok"
148    } else {
149        // claude models (sonnet, opus, haiku) and unknown → default to claude
150        "claude"
151    }
152}
153
154/// Validate that an agent name is known and installed.
155///
156/// Performs lightweight checks (known agent + binary in PATH) without
157/// a test invocation. Used by `/agent` before persisting.
158pub fn validate_agent(agent: &str) -> Result<()> {
159    let agents = crate::engine::configured_agents();
160    if !agents.iter().any(|a| a == agent) {
161        anyhow::bail!(
162            "unknown agent '{}'. Available: {}",
163            agent,
164            agents.join(", ")
165        );
166    }
167    if !crate::cmd_cache::command_exists(agent) {
168        anyhow::bail!("agent '{}' not found in PATH. Is it installed?", agent);
169    }
170    Ok(())
171}
172
173/// Validate a model spec by running a quick test invocation.
174///
175/// Sends "hello" to the agent with the specified model to verify
176/// the agent binary exists and the model is available.
177pub async fn validate_model(spec: &ModelSpec) -> Result<()> {
178    // 1. Check agent is known and installed
179    validate_agent(&spec.agent)?;
180
181    // 3. For opencode, pre-check against `opencode models` list (fast rejection)
182    if spec.agent == "opencode" {
183        if let Ok(models) = list_opencode_models().await {
184            if !models.iter().any(|m| m == &spec.model) {
185                let suggestions: Vec<_> = models
186                    .iter()
187                    .filter(|m| m.contains(&spec.model) || spec.model.contains(m.as_str()))
188                    .take(10)
189                    .cloned()
190                    .collect();
191                let shown = if suggestions.is_empty() {
192                    models.iter().take(20).cloned().collect()
193                } else {
194                    suggestions
195                };
196                anyhow::bail!(
197                    "model '{}' not found in opencode. Available models:\n{}",
198                    spec.model,
199                    shown.join("\n")
200                );
201            }
202        }
203    }
204
205    // 4. Test invocation — always run to verify the model actually works
206    eprintln!("Testing {}:{} ...", spec.agent, spec.model);
207    let result = test_invoke(&spec.agent, &spec.model).await;
208    match result {
209        Ok(_) => Ok(()),
210        Err(e) => anyhow::bail!(
211            "model '{}' on agent '{}' is not available: {e}",
212            spec.model,
213            spec.agent
214        ),
215    }
216}
217
218/// List available opencode models via `opencode models`.
219async fn list_opencode_models() -> Result<Vec<String>> {
220    match crate::engine::runner::agents::opencode::run_opencode_models_cli(
221        std::time::Duration::from_secs(10),
222    )
223    .await
224    {
225        (models, "ok" | "empty_output") => Ok(models),
226        (_, reason) => anyhow::bail!("opencode models failed: {reason}"),
227    }
228}
229
230/// Run a minimal test invocation to verify agent+model availability.
231async fn test_invoke(agent: &str, model: &str) -> Result<()> {
232    let msg = "Reply with just the word 'ok'.";
233    // Reuse invoke_agent — it already handles temp files, timeout, and error checking
234    invoke_agent(agent, model, "", msg).await?;
235    Ok(())
236}
237
238/// Assembled context for a chat invocation.
239pub struct ChatContext {
240    /// System prompt (instructions, role, commands)
241    pub system: String,
242    /// User-visible context prepended to the message (live state, memories)
243    pub state: String,
244}
245
246/// Assemble the chat context from SQLite state and live system info.
247pub async fn assemble_context(store: &TaskStore, session_id: &str) -> Result<ChatContext> {
248    // 1. Gather memories from KV (keys matching control:memory:{session}:*)
249    // Use kv_list_prefix which correctly escapes LIKE metacharacters so that
250    // session names containing '_' or '%' don't leak memories from other sessions.
251    let memory_prefix = format!("control:memory:{session_id}:");
252    let raw_memories = store
253        .kv_list_prefix(&memory_prefix)
254        .await
255        .unwrap_or_else(|e| {
256            tracing::warn!(session_id, error = %e, "failed to load control memories from store");
257            Vec::new()
258        });
259    let memories = if raw_memories.is_empty() {
260        "(none)".to_string()
261    } else {
262        raw_memories
263            .iter()
264            .map(|(key, value)| {
265                let short_key = key.strip_prefix("control:memory:").unwrap_or(key);
266                format!("- **{short_key}**: {value}")
267            })
268            .collect::<Vec<_>>()
269            .join("\n")
270    };
271
272    // 2. Recent conversation summaries
273    let summaries = store
274        .control_recent_summaries(session_id, 20)
275        .await
276        .unwrap_or_else(|e| {
277            tracing::warn!(session_id, error = %e, "failed to load control summaries from store");
278            Vec::new()
279        });
280    let recent_summaries = if summaries.is_empty() {
281        "(no recent conversation)".to_string()
282    } else {
283        summaries
284            .iter()
285            .map(|s| format!("- {s}"))
286            .collect::<Vec<_>>()
287            .join("\n")
288    };
289
290    // 3-5. Run independent subprocess calls in parallel (best-effort, 10s timeout each)
291    let (service_status, task_list, job_list) = tokio::join!(
292        async {
293            let mut cmd = tokio::process::Command::new("brew");
294            cmd.args(["services", "info", "orch"]).kill_on_drop(true);
295            match timeout(SUBPROCESS_TIMEOUT, cmd.output()).await {
296                Ok(Ok(output)) if output.status.success() => {
297                    String::from_utf8_lossy(&output.stdout).trim().to_string()
298                }
299                _ => "(could not check service status)".to_string(),
300            }
301        },
302        async {
303            let mut cmd = tokio::process::Command::new("orch");
304            cmd.args(["task", "list"]).kill_on_drop(true);
305            match timeout(SUBPROCESS_TIMEOUT, cmd.output()).await {
306                Ok(Ok(output)) if output.status.success() => {
307                    String::from_utf8_lossy(&output.stdout).trim().to_string()
308                }
309                _ => "(could not fetch live state)".to_string(),
310            }
311        },
312        async {
313            let mut cmd = tokio::process::Command::new("orch");
314            cmd.args(["job", "list"]).kill_on_drop(true);
315            match timeout(SUBPROCESS_TIMEOUT, cmd.output()).await {
316                Ok(Ok(output)) if output.status.success() => {
317                    String::from_utf8_lossy(&output.stdout).trim().to_string()
318                }
319                _ => "(no jobs configured)".to_string(),
320            }
321        },
322    );
323
324    let version = env!("ORCH_VERSION");
325    let state = format!(
326        "## 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}"
327    );
328
329    // 6. Build system prompt (instructions only, no live state in system prompt)
330    let system = SYSTEM_TEMPLATE.to_string();
331
332    Ok(ChatContext { system, state })
333}
334
335/// Get the current model spec from KV, defaulting to `"sonnet"` when the key is absent.
336/// Returns an error if the store read fails.
337pub async fn get_model(store: &TaskStore) -> Result<String> {
338    Ok(store
339        .kv_get("control:model")
340        .await
341        .context("reading control:model from store")?
342        .unwrap_or_else(|| "sonnet".to_string()))
343}
344
345/// Get the current agent from KV, defaulting to `"claude"` when the key is absent.
346/// Returns an error if the store read fails.
347pub async fn get_agent(store: &TaskStore) -> Result<String> {
348    Ok(store
349        .kv_get("control:agent")
350        .await
351        .context("reading control:agent from store")?
352        .unwrap_or_else(|| "claude".to_string()))
353}
354
355/// Get the current control session spec from KV.
356pub async fn get_current_spec(store: &TaskStore) -> Result<ModelSpec> {
357    Ok(ModelSpec {
358        agent: get_agent(store).await?,
359        model: get_model(store).await?,
360    })
361}
362
363/// Set the current model in KV.
364pub async fn set_model(store: &TaskStore, model: &str) -> Result<()> {
365    store
366        .kv_set("control:model", model)
367        .await
368        .context("setting control model")?;
369    Ok(())
370}
371
372/// Set the current agent in KV.
373pub async fn set_agent(store: &TaskStore, agent: &str) -> Result<()> {
374    store
375        .kv_set("control:agent", agent)
376        .await
377        .context("setting control agent")?;
378    Ok(())
379}
380
381/// Set the current model and agent in KV.
382pub async fn set_model_spec(store: &TaskStore, spec: &ModelSpec) -> Result<()> {
383    set_model(store, &spec.model).await?;
384    set_agent(store, &spec.agent).await?;
385    Ok(())
386}
387
388/// Set the current agent and its default model in KV.
389pub async fn set_agent_spec(store: &TaskStore, agent: &str) -> Result<ModelSpec> {
390    let spec = ModelSpec {
391        agent: agent.to_string(),
392        model: default_model_for_agent(agent).to_string(),
393    };
394    set_model_spec(store, &spec).await?;
395    Ok(spec)
396}
397
398fn display_model(model: &str) -> &str {
399    if model.is_empty() {
400        "default"
401    } else {
402        model
403    }
404}
405
406fn format_current_spec(spec: &ModelSpec) -> String {
407    format!("Current: {}:{}", spec.agent, display_model(&spec.model))
408}
409
410fn format_switched_spec(spec: &ModelSpec) -> String {
411    format!(
412        "Switched to **{}:{}** (validated)",
413        spec.agent,
414        display_model(&spec.model)
415    )
416}
417
418async fn handle_control_command(
419    store: &TaskStore,
420    session_id: &str,
421    channel_thread: Option<&str>,
422    command: ControlCommand,
423) -> Result<String> {
424    match command {
425        ControlCommand::Model(None) => Ok(format_current_spec(&get_current_spec(store).await?)),
426        ControlCommand::Model(Some(new_spec)) => {
427            let spec = parse_model_spec(&new_spec);
428            validate_model(&spec).await?;
429            set_model_spec(store, &spec).await?;
430            // Reset session UUID so next message starts a fresh conversation
431            reset_session_uuid(store, session_id).await?;
432            Ok(format_switched_spec(&spec))
433        }
434        ControlCommand::Agent(None) => Ok(format_current_spec(&get_current_spec(store).await?)),
435        ControlCommand::Agent(Some(agent)) => {
436            validate_agent(&agent)?;
437            let spec = set_agent_spec(store, &agent).await?;
438            // Reset session UUID so next message starts a fresh conversation
439            reset_session_uuid(store, session_id).await?;
440            Ok(format_switched_spec(&spec))
441        }
442        ControlCommand::Restart => {
443            // Store the conversation key so the post-restart startup check can send a confirmation.
444            if let Some(ct) = channel_thread {
445                store.kv_set("control:restart_pending", ct).await.ok();
446            }
447            // Delay the SIGTERM slightly so the acknowledgment reply is delivered first.
448            tokio::spawn(async {
449                tokio::time::sleep(Duration::from_secs(2)).await;
450                let pid = std::process::id().to_string();
451                let _ = tokio::process::Command::new("kill")
452                    .args(["-TERM", &pid])
453                    .status()
454                    .await;
455            });
456            Ok("Restarting orch service\u{2026} I'll confirm here when back online.".to_string())
457        }
458    }
459}
460
461/// Handle a control-session command if present.
462pub async fn maybe_handle_control_command(
463    store: &TaskStore,
464    session_id: &str,
465    channel_thread: Option<&str>,
466    message: &str,
467) -> Result<Option<String>> {
468    Ok(match parse_control_command(message) {
469        Some(command) => {
470            Some(handle_control_command(store, session_id, channel_thread, command).await?)
471        }
472        None => None,
473    })
474}
475
476/// Memory entry parsed from an agent response.
477#[derive(Debug, Clone, PartialEq, Eq)]
478pub struct MemoryEntry {
479    pub key: String,
480    pub value: String,
481}
482
483/// Parsed agent response metadata.
484#[derive(Debug, Clone, PartialEq, Eq)]
485pub struct ParsedResponse {
486    pub clean_text: String,
487    pub summary: Option<String>,
488    pub memories: Vec<MemoryEntry>,
489}
490
491static MEMORY_RE: LazyLock<Regex> = LazyLock::new(|| {
492    Regex::new(r#"(?s)<memory\s+key="([^"]+)">(.+?)</memory>"#).expect("valid memory regex")
493});
494
495static SUMMARY_RE: LazyLock<Regex> =
496    LazyLock::new(|| Regex::new(r#"(?s)<summary>(.+?)</summary>"#).expect("valid summary regex"));
497
498/// Parse an agent response, extracting optional `<summary>` and `<memory>` tags.
499pub fn parse_response(raw: &str) -> ParsedResponse {
500    let mut clean = raw.to_string();
501    let memories = MEMORY_RE
502        .captures_iter(raw)
503        .map(|caps| MemoryEntry {
504            key: caps[1].trim().to_string(),
505            value: caps[2].trim().to_string(),
506        })
507        .collect::<Vec<_>>();
508    clean = MEMORY_RE.replace_all(&clean, "").to_string();
509
510    let summary = SUMMARY_RE
511        .captures(raw)
512        .map(|caps| caps[1].trim().to_string());
513    clean = SUMMARY_RE.replace_all(&clean, "").to_string();
514
515    ParsedResponse {
516        clean_text: clean.trim().to_string(),
517        summary,
518        memories,
519    }
520}
521
522async fn persist_memories(
523    store: &TaskStore,
524    session_id: &str,
525    memories: &[MemoryEntry],
526) -> Result<()> {
527    for memory in memories {
528        store
529            .kv_set(
530                &format!("control:memory:{session_id}:{}", memory.key),
531                &memory.value,
532            )
533            .await
534            .with_context(|| format!("storing control memory {}", memory.key))?;
535    }
536    Ok(())
537}
538
539/// Result of an agent invocation, including response text and token usage.
540pub struct InvokeResult {
541    pub text: String,
542    pub input_tokens: Option<u64>,
543    pub output_tokens: Option<u64>,
544}
545
546fn token_u64_to_i64_saturating(tokens: u64) -> i64 {
547    i64::try_from(tokens).unwrap_or(i64::MAX)
548}
549
550fn opt_token_u64_to_i64_saturating(tokens: Option<u64>) -> Option<i64> {
551    tokens.map(token_u64_to_i64_saturating)
552}
553
554fn total_tokens_u64_to_i64_saturating(
555    input_tokens: Option<u64>,
556    output_tokens: Option<u64>,
557) -> Option<i64> {
558    match (input_tokens, output_tokens) {
559        (Some(input), Some(output)) => {
560            Some(token_u64_to_i64_saturating(input.saturating_add(output)))
561        }
562        (Some(tokens), None) | (None, Some(tokens)) => Some(token_u64_to_i64_saturating(tokens)),
563        _ => None,
564    }
565}
566
567/// Prepare temp directory for agent invocation files.
568///
569/// Uses `~/.orch/state/control/{pid}/` for isolation between concurrent processes.
570async fn prepare_temp_dir() -> Result<String> {
571    let home = crate::home::orch_home()?;
572    let dir = format!("{}/state/control/{}", home.display(), std::process::id());
573    tokio::fs::create_dir_all(&dir).await?;
574    Ok(dir)
575}
576
577/// Invoke an agent CLI one-shot and return structured result with token usage.
578///
579/// Delegates to `runner::direct::run_direct()` for command building, execution,
580/// parsing, and error classification.
581///
582/// Writes temp files to `~/.orch/state/control/{pid}/`.
583pub async fn invoke_agent(
584    agent: &str,
585    model: &str,
586    context: &str,
587    message: &str,
588) -> Result<InvokeResult> {
589    let dir = prepare_temp_dir().await?;
590
591    let result = crate::engine::runner::direct::run_direct(
592        agent,
593        Some(model),
594        context,
595        message,
596        AGENT_TIMEOUT,
597        &dir,
598    )
599    .await?;
600
601    Ok(InvokeResult {
602        text: result.text,
603        input_tokens: result.input_tokens,
604        output_tokens: result.output_tokens,
605    })
606}
607
608/// KV key for storing the session UUID for a named session.
609fn session_uuid_key(session_name: &str) -> String {
610    format!("control:session_id:{session_name}")
611}
612
613/// Get or create a session UUID for the given session name.
614///
615/// On first call for a session, generates a new UUID and persists it to KV.
616/// Subsequent calls return the same UUID, giving agents conversation continuity
617/// via `--session-id`.
618/// Get or create a session UUID for the given session name.
619///
620/// On first call for a session, generates a new UUID and persists it to KV.
621/// Subsequent calls return the same UUID, giving agents conversation continuity
622/// via `--session-id` (new) or `--resume` (existing).
623///
624/// Returns `(uuid, is_new)` so callers know whether to create or resume the session.
625pub async fn get_or_create_session_uuid(
626    store: &TaskStore,
627    session_name: &str,
628) -> Result<(String, bool)> {
629    let key = session_uuid_key(session_name);
630    let candidate = Uuid::new_v4().to_string();
631    let stored = store
632        .kv_insert_if_absent(&key, &candidate)
633        .await
634        .context("storing session UUID")?;
635    let is_new = stored == candidate;
636    if is_new {
637        tracing::info!(session_name, uuid = %stored, "created new chat session UUID");
638    }
639    Ok((stored, is_new))
640}
641
642/// Reset the session UUID so the next message starts a fresh conversation.
643///
644/// Called when agent or model changes, since the new agent won't share
645/// the previous agent's session state. Deletes the stored UUID so the next
646/// `get_or_create_session_uuid` call generates a fresh one.
647pub async fn reset_session_uuid(store: &TaskStore, session_name: &str) -> Result<()> {
648    let key = session_uuid_key(session_name);
649    store
650        .kv_delete(&key)
651        .await
652        .context("resetting session UUID")?;
653    tracing::info!(session_name, "reset chat session UUID");
654    Ok(())
655}
656
657/// Invoke an agent one-shot with `--session-id` or `--resume` for conversation continuity.
658///
659/// For claude-compatible agents (claude, kimi, minimax), passes `--session-id <uuid>`
660/// when creating a new session, or `--resume <uuid>` when continuing an existing one.
661/// Other agents get a plain one-shot call.
662pub async fn invoke_agent_with_session(
663    agent: &str,
664    model: &str,
665    context: &str,
666    message: &str,
667    session_uuid: Option<&str>,
668    resume: bool,
669) -> Result<InvokeResult> {
670    let dir = prepare_temp_dir().await?;
671
672    let session = session_uuid.map(|sid| crate::engine::runner::direct::SessionContinuation {
673        session_id: sid,
674        resume,
675    });
676
677    let result = crate::engine::runner::direct::run_direct_with_session(
678        agent,
679        Some(model),
680        context,
681        message,
682        AGENT_TIMEOUT,
683        &dir,
684        session,
685    )
686    .await?;
687
688    Ok(InvokeResult {
689        text: result.text,
690        input_tokens: result.input_tokens,
691        output_tokens: result.output_tokens,
692    })
693}
694
695/// High-level entry point: process a user message and return the assistant response.
696///
697/// Handles `/model` and `/agent` commands, assembles context, invokes the agent,
698/// parses the response for summaries, and stores both user and assistant messages.
699pub async fn send_message(
700    store: &TaskStore,
701    session_id: &str,
702    channel: &str,
703    channel_thread: Option<&str>,
704    message: &str,
705) -> Result<String> {
706    // Handle control commands — don't store as messages.
707    if let Some(response) =
708        maybe_handle_control_command(store, session_id, channel_thread, message).await?
709    {
710        return Ok(response);
711    }
712
713    // Acquire per-session lock — serializes concurrent invocations for the same session.
714    // Messages that arrive while an invocation is running will queue here and execute in order.
715    let session_lock = {
716        let mut map = SESSION_LOCKS.lock().unwrap_or_else(|e| e.into_inner());
717        Arc::clone(
718            map.entry(session_id.to_string())
719                .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))),
720        )
721    };
722    let _guard = session_lock.lock().await;
723
724    // Store user message
725    store
726        .insert_control_message(
727            session_id,
728            "user",
729            channel,
730            channel_thread,
731            message,
732            None,
733            None,
734            None,
735            None,
736            None,
737            None,
738            None,
739        )
740        .await?;
741
742    // Assemble context
743    let ctx = assemble_context(store, session_id).await?;
744
745    // Resolve model and agent from KV
746    let model = get_model(store).await?;
747    let agent = get_agent(store).await?;
748
749    // Prepend live state to user message so the LLM sees it directly
750    let full_message = format!("{}\n\n---\n\n{}", ctx.state, message);
751
752    // Get or create a session UUID for conversation continuity
753    let (session_uuid, is_new) = get_or_create_session_uuid(store, session_id).await?;
754
755    // When the combined payload exceeds the per-turn limit, split it across multiple
756    // sequential turns so the user's question (trailing content) is never truncated.
757    // For session-supporting agents (claude, kimi, minimax), use --resume for continuity
758    // and persist ALL intermediate assistant responses so local transcript matches provider.
759    // For other agents (codex, opencode), split into chunks and send each as independent
760    // invocation (no --resume), each getting its own persisted assistant response.
761    let agent_has_sessions = crate::engine::runner::agents::supports_session(&agent);
762    let (chunks, delivery_mode) = if full_message.len() > MAX_CONTROL_MESSAGE_BYTES {
763        tracing::info!(
764            bytes = full_message.len(),
765            limit = MAX_CONTROL_MESSAGE_BYTES,
766            agent = %agent,
767            agent_has_sessions,
768            "control session payload exceeds limit"
769        );
770        if agent_has_sessions {
771            // Session agents: split into chunks for sequential delivery with --resume
772            (
773                build_control_chunks(&ctx.state, message, true),
774                DeliveryMode::Chunked,
775            )
776        } else {
777            // Non-session agents: split into chunks, send each independently
778            (
779                build_control_chunks(&ctx.state, message, false),
780                DeliveryMode::ChunkedNonSession,
781            )
782        }
783    } else {
784        (vec![full_message], DeliveryMode::Single)
785    };
786
787    let _num_chunks = chunks.len();
788
789    // Deliver chunks based on mode
790    let result = match delivery_mode {
791        DeliveryMode::Chunked => {
792            // Session-based delivery (claude, kimi, minimax) — use --resume for continuity
793            // Persist ALL intermediate responses so local DB matches provider transcript.
794            let mut active_uuid = session_uuid.clone();
795            let mut is_first_chunk = !is_new;
796            let mut last_result: Option<InvokeResult> = None;
797
798            for chunk in chunks.iter() {
799                let r = invoke_agent_with_session(
800                    &agent,
801                    &model,
802                    &ctx.system,
803                    chunk,
804                    Some(&active_uuid),
805                    is_first_chunk,
806                )
807                .await;
808                let chunk_result = match r {
809                    Ok(result) => result,
810                    Err(e) => {
811                        let err_str = e.to_string();
812                        let is_stale = err_str.contains("stale session");
813                        let is_timeout = err_str.contains("timeout after")
814                            || err_str.contains("timed out after");
815                        if is_stale || is_timeout {
816                            if is_stale {
817                                tracing::info!(
818                                    session_id,
819                                    "chat session UUID expired; resetting and retrying with fresh session"
820                                );
821                            } else {
822                                tracing::info!(
823                                    session_id,
824                                    "chat session timed out; resetting and retrying with fresh session"
825                                );
826                            }
827                            reset_session_uuid(store, session_id).await?;
828                            let (fresh_uuid, _) =
829                                get_or_create_session_uuid(store, session_id).await?;
830                            active_uuid = fresh_uuid;
831                            is_first_chunk = false;
832                            let fresh_result = invoke_agent_with_session(
833                                &agent,
834                                &model,
835                                &ctx.system,
836                                chunk,
837                                Some(&active_uuid),
838                                false,
839                            )
840                            .await?;
841                            let result = if is_timeout {
842                                InvokeResult {
843                                    text: format!(
844                                        "_(New session started.)_\n\n{}",
845                                        fresh_result.text
846                                    ),
847                                    ..fresh_result
848                                }
849                            } else {
850                                fresh_result
851                            };
852                            result
853                        } else {
854                            return Err(e);
855                        }
856                    }
857                };
858
859                // Persist this chunk's assistant response (even intermediate ones)
860                let parsed = parse_response(&chunk_result.text);
861                let total_tokens = total_tokens_u64_to_i64_saturating(
862                    chunk_result.input_tokens,
863                    chunk_result.output_tokens,
864                );
865                let input_tokens = opt_token_u64_to_i64_saturating(chunk_result.input_tokens);
866                let output_tokens = opt_token_u64_to_i64_saturating(chunk_result.output_tokens);
867                let cost_usd = match (chunk_result.input_tokens, chunk_result.output_tokens) {
868                    (Some(input), Some(output)) => {
869                        let pricing = crate::store::pricing_for_model(&model);
870                        let usage = crate::store::TokenUsage {
871                            input_tokens: input,
872                            output_tokens: output,
873                        };
874                        Some(pricing.estimate_cost_usd(usage).total_cost_usd)
875                    }
876                    _ => None,
877                };
878                store
879                    .insert_control_message(
880                        session_id,
881                        "assistant",
882                        channel,
883                        channel_thread,
884                        &parsed.clean_text,
885                        parsed.summary.as_deref(),
886                        Some(&model),
887                        Some(&agent),
888                        input_tokens,
889                        output_tokens,
890                        total_tokens,
891                        cost_usd,
892                    )
893                    .await?;
894                persist_memories(store, session_id, &parsed.memories).await?;
895
896                last_result = Some(chunk_result);
897            }
898
899            // Return final clean_text directly — all chunks already persisted inline
900            let final_parsed = parse_response(&last_result.expect("at least one chunk").text);
901            final_parsed.clean_text
902        }
903        DeliveryMode::ChunkedNonSession => {
904            // Non-session agents (codex, opencode) — each chunk is independent invocation
905            // No --resume, but each chunk's response is persisted to keep transcript complete.
906            let mut last_text = String::new();
907
908            for chunk in &chunks {
909                let chunk_result = invoke_agent(&agent, &model, &ctx.system, chunk).await?;
910
911                // Persist this chunk's assistant response
912                let parsed = parse_response(&chunk_result.text);
913                let total_tokens = total_tokens_u64_to_i64_saturating(
914                    chunk_result.input_tokens,
915                    chunk_result.output_tokens,
916                );
917                let input_tokens = opt_token_u64_to_i64_saturating(chunk_result.input_tokens);
918                let output_tokens = opt_token_u64_to_i64_saturating(chunk_result.output_tokens);
919                let cost_usd = match (chunk_result.input_tokens, chunk_result.output_tokens) {
920                    (Some(input), Some(output)) => {
921                        let pricing = crate::store::pricing_for_model(&model);
922                        let usage = crate::store::TokenUsage {
923                            input_tokens: input,
924                            output_tokens: output,
925                        };
926                        Some(pricing.estimate_cost_usd(usage).total_cost_usd)
927                    }
928                    _ => None,
929                };
930                store
931                    .insert_control_message(
932                        session_id,
933                        "assistant",
934                        channel,
935                        channel_thread,
936                        &parsed.clean_text,
937                        parsed.summary.as_deref(),
938                        Some(&model),
939                        Some(&agent),
940                        input_tokens,
941                        output_tokens,
942                        total_tokens,
943                        cost_usd,
944                    )
945                    .await?;
946                persist_memories(store, session_id, &parsed.memories).await?;
947
948                last_text = parsed.clean_text;
949            }
950
951            last_text
952        }
953        DeliveryMode::Single => {
954            // Single message delivery (under-limit messages) — use invoke_agent_with_session
955            // for conversation continuity with --session-id / --resume
956            let mut active_uuid = session_uuid.clone();
957            let mut is_first_chunk = !is_new;
958            let chunk_result = loop {
959                let r = invoke_agent_with_session(
960                    &agent,
961                    &model,
962                    &ctx.system,
963                    &chunks[0],
964                    Some(&active_uuid),
965                    is_first_chunk,
966                )
967                .await;
968                match r {
969                    Ok(result) => break result,
970                    Err(e) => {
971                        let err_str = e.to_string();
972                        let is_stale = err_str.contains("stale session");
973                        let is_timeout = err_str.contains("timeout after")
974                            || err_str.contains("timed out after");
975                        if is_stale || is_timeout {
976                            if is_stale {
977                                tracing::info!(
978                                    session_id,
979                                    "chat session UUID expired; resetting and retrying with fresh session"
980                                );
981                            } else {
982                                tracing::info!(
983                                    session_id,
984                                    "chat session timed out; resetting and retrying with fresh session"
985                                );
986                            }
987                            reset_session_uuid(store, session_id).await?;
988                            let (fresh_uuid, _) =
989                                get_or_create_session_uuid(store, session_id).await?;
990                            active_uuid = fresh_uuid;
991                            is_first_chunk = false;
992                            continue;
993                        } else {
994                            return Err(e);
995                        }
996                    }
997                }
998            };
999
1000            // Parse response for summary and memory tags
1001            let parsed = parse_response(&chunk_result.text);
1002
1003            // Store assistant message with token usage.
1004            // Must happen before persist_memories so that if the DB write fails the
1005            // KV store is not left with orphaned memory keys that have no corresponding
1006            // conversation history entry.
1007            let total_tokens = total_tokens_u64_to_i64_saturating(
1008                chunk_result.input_tokens,
1009                chunk_result.output_tokens,
1010            );
1011            let input_tokens = opt_token_u64_to_i64_saturating(chunk_result.input_tokens);
1012            let output_tokens = opt_token_u64_to_i64_saturating(chunk_result.output_tokens);
1013            let cost_usd = match (chunk_result.input_tokens, chunk_result.output_tokens) {
1014                (Some(input), Some(output)) => {
1015                    let pricing = crate::store::pricing_for_model(&model);
1016                    let usage = crate::store::TokenUsage {
1017                        input_tokens: input,
1018                        output_tokens: output,
1019                    };
1020                    Some(pricing.estimate_cost_usd(usage).total_cost_usd)
1021                }
1022                _ => None,
1023            };
1024            store
1025                .insert_control_message(
1026                    session_id,
1027                    "assistant",
1028                    channel,
1029                    channel_thread,
1030                    &parsed.clean_text,
1031                    parsed.summary.as_deref(),
1032                    Some(&model),
1033                    Some(&agent),
1034                    input_tokens,
1035                    output_tokens,
1036                    total_tokens,
1037                    cost_usd,
1038                )
1039                .await?;
1040            persist_memories(store, session_id, &parsed.memories).await?;
1041
1042            parsed.clean_text
1043        }
1044    };
1045
1046    // Release the per-session lock before eviction check.
1047    drop(_guard);
1048
1049    // Evict the session lock entry if no concurrent waiters are holding a reference.
1050    // After drop(_guard) above the remaining strong references are:
1051    // - the map entry
1052    // - the local `session_lock` clone held by this function.
1053    // That sums to 2 when there are no other waiters. If another waiter entered
1054    // concurrently the count will be > 2 and we should not evict. Only remove the
1055    // map entry when the strong_count == 2.
1056    if Arc::strong_count(&session_lock) == 2 {
1057        let mut map = SESSION_LOCKS.lock().unwrap_or_else(|e| e.into_inner());
1058        // Re-check under the map lock: a concurrent caller may have cloned the Arc
1059        // between our strong_count check above and acquiring the map lock.
1060        if map
1061            .get(session_id)
1062            .is_some_and(|l| Arc::strong_count(l) == 2)
1063        {
1064            map.remove(session_id);
1065        }
1066    }
1067
1068    Ok(result)
1069}
1070
1071/// Split `text` into chunks of at most `max_bytes` bytes each.
1072/// Prefers splitting at newline boundaries to preserve readability.
1073fn chunk_text(text: &str, max_bytes: usize) -> Vec<String> {
1074    if text.len() <= max_bytes {
1075        return vec![text.to_string()];
1076    }
1077    let mut chunks = Vec::new();
1078    let mut start = 0;
1079    let total = text.len();
1080
1081    while start < total {
1082        let mut end = (start + max_bytes).min(total);
1083        // Respect UTF-8 char boundaries
1084        while end > start && !text.is_char_boundary(end) {
1085            end -= 1;
1086        }
1087        if end == start {
1088            // Single char exceeds limit — advance past it
1089            if let Some(ch) = text[start..].chars().next() {
1090                end = start + ch.len_utf8();
1091            } else {
1092                break;
1093            }
1094        }
1095        // Prefer splitting at a newline within the chunk
1096        if end < total {
1097            if let Some(nl_offset) = text[start..end].rfind('\n') {
1098                end = start + nl_offset + 1;
1099            }
1100        }
1101        chunks.push(text[start..end].to_string());
1102        start = end;
1103    }
1104    chunks
1105}
1106
1107/// Build control session message chunks for agents with or without session support.
1108///
1109/// For session-supporting agents (claude, kimi, minimax), chunks are raw text pieces
1110/// that will be sent sequentially with `--resume` for conversation continuity.
1111///
1112/// For non-session agents (codex, opencode), each chunk is prefixed with a header
1113/// indicating its position in the sequence, since each invocation is independent.
1114fn build_control_chunks(state: &str, message: &str, agent_has_sessions: bool) -> Vec<String> {
1115    let header_reserve = if agent_has_sessions { 0 } else { 200 };
1116    let effective_limit = MAX_CONTROL_MESSAGE_BYTES.saturating_sub(header_reserve);
1117    let mut parts = Vec::new();
1118
1119    // Split state into chunks
1120    let state_chunks = chunk_text(state, effective_limit);
1121    for (i, chunk) in state_chunks.iter().enumerate() {
1122        let prefixed = if agent_has_sessions {
1123            chunk.clone()
1124        } else {
1125            format!(
1126                "[Context chunk {}/{}]\n{}",
1127                i + 1,
1128                state_chunks.len(),
1129                chunk
1130            )
1131        };
1132        parts.push(prefixed);
1133    }
1134
1135    // Split user message if needed, then add as final chunk(s)
1136    let msg_chunks = chunk_text(message, effective_limit);
1137    for (i, chunk) in msg_chunks.iter().enumerate() {
1138        let prefixed = if agent_has_sessions {
1139            chunk.clone()
1140        } else if msg_chunks.len() == 1 {
1141            format!("[Context complete. User question:]\n{}", chunk)
1142        } else {
1143            format!(
1144                "[User question chunk {}/{}]\n{}",
1145                i + 1,
1146                msg_chunks.len(),
1147                chunk
1148            )
1149        };
1150        parts.push(prefixed);
1151    }
1152    parts
1153}
1154
1155#[cfg(test)]
1156mod tests {
1157    use super::*;
1158    use crate::store::TaskStore;
1159
1160    #[test]
1161    fn parse_spec_agent_model() {
1162        let spec = parse_model_spec("minimax:sonnet");
1163        assert_eq!(spec.agent, "minimax");
1164        assert_eq!(spec.model, "sonnet");
1165    }
1166
1167    #[test]
1168    fn parse_spec_opencode_model() {
1169        let spec = parse_model_spec("opencode:minimax-m2.5-free");
1170        assert_eq!(spec.agent, "opencode");
1171        assert_eq!(spec.model, "minimax-m2.5-free");
1172    }
1173
1174    #[test]
1175    fn parse_spec_model_only_infers_claude() {
1176        let spec = parse_model_spec("sonnet");
1177        assert_eq!(spec.agent, "claude");
1178        assert_eq!(spec.model, "sonnet");
1179    }
1180
1181    #[test]
1182    fn parse_spec_model_only_infers_codex() {
1183        let spec = parse_model_spec("gpt-4o");
1184        assert_eq!(spec.agent, "codex");
1185        assert_eq!(spec.model, "gpt-4o");
1186    }
1187
1188    #[test]
1189    fn parse_spec_model_only_infers_opencode() {
1190        let spec = parse_model_spec("deepseek-r1");
1191        assert_eq!(spec.agent, "opencode");
1192        assert_eq!(spec.model, "deepseek-r1");
1193    }
1194
1195    #[tokio::test]
1196    async fn assemble_context_includes_memories_and_summaries() {
1197        let store = TaskStore::open_memory().await.unwrap();
1198        store
1199            .kv_set("control:memory:default:tz", "User timezone is UTC-3")
1200            .await
1201            .unwrap();
1202        store
1203            .insert_control_message(
1204                "default",
1205                "assistant",
1206                "cli",
1207                None,
1208                "response",
1209                Some("checked status"),
1210                Some("sonnet"),
1211                Some("claude"),
1212                None,
1213                None,
1214                None,
1215                None,
1216            )
1217            .await
1218            .unwrap();
1219
1220        let ctx = assemble_context(&store, "default").await.unwrap();
1221        assert!(ctx.state.contains("UTC-3"), "should include memories");
1222        assert!(
1223            ctx.state.contains("checked status"),
1224            "should include summaries"
1225        );
1226    }
1227
1228    #[test]
1229    fn extract_summary_from_response() {
1230        let response = "Here are your tasks.\n\n<summary>listed 3 active tasks</summary>";
1231        let parsed = parse_response(response);
1232        assert_eq!(parsed.summary.as_deref(), Some("listed 3 active tasks"));
1233        assert!(!parsed.clean_text.contains("<summary>"));
1234    }
1235
1236    #[test]
1237    fn extract_summary_missing() {
1238        let response = "No tasks found.";
1239        let parsed = parse_response(response);
1240        assert_eq!(parsed.summary, None);
1241        assert_eq!(parsed.clean_text, "No tasks found.");
1242    }
1243
1244    #[test]
1245    fn extract_summary_closing_tag_before_opening_does_not_panic() {
1246        // </summary> appears before <summary> — must not panic
1247        let response = "Example: </summary> is a closing tag.\n<summary>actual summary</summary>";
1248        let parsed = parse_response(response);
1249        assert_eq!(parsed.summary.as_deref(), Some("actual summary"));
1250        assert!(!parsed.clean_text.contains("<summary>"));
1251    }
1252
1253    #[test]
1254    fn extract_memory_tags() {
1255        let response = r#"Remember these.
1256
1257<memory key="tz">User timezone is UTC-3</memory>
1258<memory key="theme">Prefer concise answers</memory>
1259
1260<summary>stored two memories</summary>"#;
1261        let parsed = parse_response(response);
1262        assert_eq!(parsed.memories.len(), 2);
1263        assert_eq!(parsed.memories[0].key, "tz");
1264        assert_eq!(parsed.memories[0].value, "User timezone is UTC-3");
1265        assert_eq!(parsed.memories[1].key, "theme");
1266        assert_eq!(parsed.memories[1].value, "Prefer concise answers");
1267        assert_eq!(parsed.summary.as_deref(), Some("stored two memories"));
1268        assert!(!parsed.clean_text.contains("<memory"));
1269    }
1270
1271    #[tokio::test]
1272    async fn persist_memories_writes_to_kv() {
1273        let store = TaskStore::open_memory().await.unwrap();
1274        let memories = vec![MemoryEntry {
1275            key: "tz".to_string(),
1276            value: "User timezone is UTC-3".to_string(),
1277        }];
1278
1279        persist_memories(&store, "default", &memories)
1280            .await
1281            .unwrap();
1282
1283        assert_eq!(
1284            store
1285                .kv_get("control:memory:default:tz")
1286                .await
1287                .unwrap()
1288                .as_deref(),
1289            Some("User timezone is UTC-3")
1290        );
1291    }
1292
1293    #[test]
1294    fn resolve_agent_from_model() {
1295        assert_eq!(infer_agent("sonnet"), "claude");
1296        assert_eq!(infer_agent("opus"), "claude");
1297        assert_eq!(infer_agent("haiku"), "claude");
1298        assert_eq!(infer_agent("gpt-4o"), "codex");
1299        assert_eq!(infer_agent("o3"), "codex");
1300        assert_eq!(infer_agent("deepseek-r1"), "opencode");
1301        assert_eq!(infer_agent("grok"), "grok");
1302    }
1303
1304    #[tokio::test]
1305    async fn send_message_model_switch() {
1306        let store = TaskStore::open_memory().await.unwrap();
1307        let model = get_model(&store).await.unwrap();
1308        assert_eq!(model, "sonnet");
1309
1310        // /model with explicit agent:model — will fail validation (no binary in test)
1311        // so test parse_model_spec + set_model_spec directly
1312        let spec = parse_model_spec("minimax:sonnet");
1313        set_model_spec(&store, &spec).await.unwrap();
1314        assert_eq!(get_model(&store).await.unwrap(), "sonnet");
1315        assert_eq!(get_agent(&store).await.unwrap(), "minimax");
1316    }
1317
1318    #[tokio::test]
1319    async fn send_message_model_show_current() {
1320        let store = TaskStore::open_memory().await.unwrap();
1321        let response = send_message(&store, "default", "cli", None, "/model")
1322            .await
1323            .unwrap();
1324        assert!(response.contains("claude"));
1325        assert!(response.contains("sonnet"));
1326    }
1327
1328    #[tokio::test]
1329    async fn send_message_agent_show_current() {
1330        let store = TaskStore::open_memory().await.unwrap();
1331        let response = send_message(&store, "default", "cli", None, "/agent")
1332            .await
1333            .unwrap();
1334        assert!(response.contains("claude"));
1335        assert!(response.contains("sonnet"));
1336    }
1337
1338    #[tokio::test]
1339    async fn send_message_agent_switches_with_default_model() {
1340        let store = TaskStore::open_memory().await.unwrap();
1341        // Directly set agent spec to bypass validation (no binary in test env)
1342        let spec = set_agent_spec(&store, "codex").await.unwrap();
1343        assert_eq!(spec.agent, "codex");
1344        assert_eq!(spec.model, "gpt-4o");
1345        assert_eq!(get_agent(&store).await.unwrap(), "codex");
1346        assert_eq!(get_model(&store).await.unwrap(), "gpt-4o");
1347    }
1348
1349    #[test]
1350    fn parse_control_command_model_and_agent() {
1351        assert_eq!(
1352            parse_control_command("/model"),
1353            Some(ControlCommand::Model(None))
1354        );
1355        assert_eq!(
1356            parse_control_command("/agent"),
1357            Some(ControlCommand::Agent(None))
1358        );
1359        assert_eq!(
1360            parse_control_command("/agent codex"),
1361            Some(ControlCommand::Agent(Some("codex".to_string())))
1362        );
1363        assert_eq!(
1364            parse_control_command("/model opencode:deepseek-r1"),
1365            Some(ControlCommand::Model(Some(
1366                "opencode:deepseek-r1".to_string()
1367            )))
1368        );
1369        assert_eq!(
1370            parse_control_command("/restart"),
1371            Some(ControlCommand::Restart)
1372        );
1373        assert_eq!(parse_control_command("/kill"), None);
1374        assert_eq!(parse_control_command("/session"), None);
1375    }
1376
1377    #[test]
1378    fn control_prompt_slash_commands_match_supported_surface() {
1379        let documented = Regex::new(r"`/([a-z]+)(?:[^`]*)`")
1380            .unwrap()
1381            .captures_iter(SYSTEM_TEMPLATE)
1382            .map(|caps| caps[1].to_string())
1383            .collect::<std::collections::BTreeSet<_>>();
1384
1385        let expected = ["agent", "model", "restart"]
1386            .into_iter()
1387            .map(str::to_string)
1388            .collect::<std::collections::BTreeSet<_>>();
1389
1390        assert_eq!(documented, expected);
1391    }
1392
1393    #[tokio::test]
1394    async fn set_agent_spec_persists_default_model() {
1395        let store = TaskStore::open_memory().await.unwrap();
1396        let spec = set_agent_spec(&store, "codex").await.unwrap();
1397        assert_eq!(spec.agent, "codex");
1398        assert_eq!(spec.model, "gpt-4o");
1399        assert_eq!(get_agent(&store).await.unwrap(), "codex");
1400        assert_eq!(get_model(&store).await.unwrap(), "gpt-4o");
1401    }
1402
1403    #[tokio::test]
1404    async fn send_message_model_switch_does_not_store() {
1405        let store = TaskStore::open_memory().await.unwrap();
1406        // set directly to avoid validation in test env
1407        let spec = parse_model_spec("claude:haiku");
1408        set_model_spec(&store, &spec).await.unwrap();
1409        let messages = store
1410            .list_control_messages("default", None, 10)
1411            .await
1412            .unwrap();
1413        assert_eq!(messages.len(), 0);
1414    }
1415
1416    #[test]
1417    fn validate_unknown_agent_fails() {
1418        let spec = ModelSpec {
1419            agent: "nonexistent".to_string(),
1420            model: "test".to_string(),
1421        };
1422        let rt = tokio::runtime::Runtime::new().unwrap();
1423        let result = rt.block_on(validate_model(&spec));
1424        assert!(result.is_err());
1425        assert!(result.unwrap_err().to_string().contains("unknown agent"));
1426    }
1427
1428    #[test]
1429    fn validate_agent_rejects_unknown() {
1430        let result = validate_agent("nonexistent");
1431        assert!(result.is_err());
1432        let err = result.unwrap_err().to_string();
1433        assert!(err.contains("unknown agent"), "got: {err}");
1434    }
1435
1436    #[test]
1437    fn validate_agent_rejects_empty() {
1438        let result = validate_agent("");
1439        assert!(result.is_err());
1440        assert!(result.unwrap_err().to_string().contains("unknown agent"));
1441    }
1442
1443    #[tokio::test]
1444    async fn agent_command_rejects_unknown_before_persisting() {
1445        let store = TaskStore::open_memory().await.unwrap();
1446        // Set a known-good agent first
1447        set_agent(&store, "claude").await.unwrap();
1448
1449        // Try switching to an unknown agent — should fail
1450        let result =
1451            maybe_handle_control_command(&store, "default", None, "/agent bogusagent").await;
1452        assert!(result.is_err(), "should reject unknown agent");
1453        assert!(result.unwrap_err().to_string().contains("unknown agent"));
1454
1455        // Verify the original agent is still set (not overwritten)
1456        assert_eq!(get_agent(&store).await.unwrap(), "claude");
1457    }
1458
1459    #[test]
1460    fn cost_usd_computed_from_token_counts() {
1461        use crate::store::{pricing_for_model, TokenUsage};
1462
1463        let model = "sonnet";
1464        let input_tokens: u64 = 1_000_000;
1465        let output_tokens: u64 = 1_000_000;
1466        let pricing = pricing_for_model(model);
1467        let usage = TokenUsage {
1468            input_tokens,
1469            output_tokens,
1470        };
1471        let cost = pricing.estimate_cost_usd(usage);
1472
1473        // sonnet: $3/M input, $15/M output → $18 total for 1M each
1474        assert!(cost.total_cost_usd > 0.0, "cost must be non-zero");
1475        assert!(
1476            (cost.total_cost_usd - 18.0).abs() < 0.01,
1477            "expected ~$18 for 1M+1M sonnet tokens"
1478        );
1479    }
1480
1481    #[test]
1482    fn token_u64_to_i64_saturating_caps_large_values() {
1483        assert_eq!(token_u64_to_i64_saturating(u64::MAX), i64::MAX);
1484    }
1485
1486    #[test]
1487    fn total_tokens_u64_to_i64_saturating_caps_single_large_value() {
1488        assert_eq!(
1489            total_tokens_u64_to_i64_saturating(Some(u64::MAX), None),
1490            Some(i64::MAX)
1491        );
1492    }
1493
1494    #[test]
1495    fn total_tokens_u64_to_i64_saturating_caps_large_sum() {
1496        assert_eq!(
1497            total_tokens_u64_to_i64_saturating(Some(u64::MAX), Some(1)),
1498            Some(i64::MAX)
1499        );
1500    }
1501
1502    #[tokio::test]
1503    async fn insert_control_message_stores_cost_usd() {
1504        use crate::store::{pricing_for_model, TokenUsage};
1505
1506        let store = crate::store::TaskStore::open_memory().await.unwrap();
1507
1508        let input_tokens: u64 = 500_000;
1509        let output_tokens: u64 = 100_000;
1510        let pricing = pricing_for_model("sonnet");
1511        let usage = TokenUsage {
1512            input_tokens,
1513            output_tokens,
1514        };
1515        let cost_usd = Some(pricing.estimate_cost_usd(usage).total_cost_usd);
1516
1517        store
1518            .insert_control_message(
1519                "default",
1520                "assistant",
1521                "cli",
1522                None,
1523                "hello",
1524                Some("greeted"),
1525                Some("sonnet"),
1526                Some("claude"),
1527                opt_token_u64_to_i64_saturating(Some(input_tokens)),
1528                opt_token_u64_to_i64_saturating(Some(output_tokens)),
1529                total_tokens_u64_to_i64_saturating(Some(input_tokens), Some(output_tokens)),
1530                cost_usd,
1531            )
1532            .await
1533            .unwrap();
1534
1535        let messages = store
1536            .list_control_messages("default", None, 10)
1537            .await
1538            .unwrap();
1539        assert_eq!(messages.len(), 1);
1540        let msg = &messages[0];
1541        assert!(msg.cost_usd.is_some(), "cost_usd must be stored, not NULL");
1542        let stored = msg.cost_usd.unwrap();
1543        assert!(stored > 0.0, "stored cost must be positive");
1544    }
1545
1546    #[tokio::test]
1547    async fn session_uuid_created_and_stable() {
1548        let store = TaskStore::open_memory().await.unwrap();
1549        let (uuid1, is_new1) = get_or_create_session_uuid(&store, "default").await.unwrap();
1550        let (uuid2, is_new2) = get_or_create_session_uuid(&store, "default").await.unwrap();
1551        assert_eq!(uuid1, uuid2, "same session should return same UUID");
1552        assert!(is_new1, "first call should be new");
1553        assert!(!is_new2, "second call should be existing");
1554        assert_eq!(uuid1.len(), 36, "should be a valid UUID string");
1555    }
1556
1557    #[tokio::test]
1558    async fn session_uuid_differs_per_session() {
1559        let store = TaskStore::open_memory().await.unwrap();
1560        let (uuid_default, _) = get_or_create_session_uuid(&store, "default").await.unwrap();
1561        let (uuid_ops, _) = get_or_create_session_uuid(&store, "ops").await.unwrap();
1562        assert_ne!(uuid_default, uuid_ops);
1563    }
1564
1565    #[tokio::test]
1566    async fn session_uuid_reset_creates_new() {
1567        let store = TaskStore::open_memory().await.unwrap();
1568        let (uuid1, _) = get_or_create_session_uuid(&store, "default").await.unwrap();
1569        reset_session_uuid(&store, "default").await.unwrap();
1570        let (uuid2, is_new2) = get_or_create_session_uuid(&store, "default").await.unwrap();
1571        assert_ne!(uuid1, uuid2, "reset should produce a new UUID");
1572        assert!(is_new2, "after reset should be new");
1573    }
1574
1575    /// Verify that the timeout error patterns used in send_message recovery match
1576    /// the actual error strings produced by the agent runner (AgentError::Timeout
1577    /// formats as "timeout after Xs"; run_shell_command formats as "agent timed out
1578    /// after Xs"). Both must trigger session reset rather than bubbling up as errors.
1579    #[test]
1580    fn timeout_error_patterns_match_runner_output() {
1581        let patterns = ["timeout after 1800s", "agent timed out after 120s"];
1582        for &msg in &patterns {
1583            let is_timeout = msg.contains("timeout after") || msg.contains("timed out after");
1584            assert!(is_timeout, "pattern should be detected as timeout: {msg:?}");
1585        }
1586        // Unrelated errors must not be misidentified as timeouts.
1587        let non_timeout = [
1588            "rate limit: 429",
1589            "stale session (session expired): abc",
1590            "auth error: invalid key",
1591        ];
1592        for &msg in &non_timeout {
1593            let is_timeout = msg.contains("timeout after") || msg.contains("timed out after");
1594            assert!(
1595                !is_timeout,
1596                "pattern must not be detected as timeout: {msg:?}"
1597            );
1598        }
1599    }
1600
1601    // ── chunk_text ──────────────────────────────────────────────────────────────
1602
1603    #[test]
1604    fn chunk_text_under_limit_returns_single_chunk() {
1605        let text = "hello world";
1606        let chunks = chunk_text(text, 100);
1607        assert_eq!(chunks, vec!["hello world".to_string()]);
1608    }
1609
1610    #[test]
1611    fn chunk_text_exact_limit_returns_single_chunk() {
1612        let text = "abcde";
1613        let chunks = chunk_text(text, 5);
1614        assert_eq!(chunks.len(), 1);
1615        assert_eq!(chunks[0], "abcde");
1616    }
1617
1618    #[test]
1619    fn chunk_text_splits_and_reassembles() {
1620        let text = "line1\nline2\nline3\nline4\nline5";
1621        let chunks = chunk_text(text, 12);
1622        let reassembled: String = chunks.join("");
1623        assert_eq!(reassembled, text);
1624    }
1625
1626    #[test]
1627    fn chunk_text_respects_utf8_boundaries() {
1628        // 4-byte emoji placed at the split boundary
1629        let content = "x".repeat(97) + "😀"; // 101 bytes; limit=100 cuts inside emoji
1630        let chunks = chunk_text(&content, 100);
1631        let reassembled: String = chunks.join("");
1632        assert_eq!(reassembled, content);
1633        for c in &chunks {
1634            assert!(
1635                std::str::from_utf8(c.as_bytes()).is_ok(),
1636                "chunk is not valid UTF-8"
1637            );
1638        }
1639    }
1640
1641    #[test]
1642    fn chunk_text_prefers_newline_boundary() {
1643        // "aaaa\nbbbb" with limit=6 — should split after the '\n', not mid-word
1644        let text = "aaaa\nbbbb";
1645        let chunks = chunk_text(text, 6);
1646        // First chunk should end with '\n' (split at newline, not byte 6 into "aaaa\nb")
1647        assert!(
1648            chunks[0].ends_with('\n'),
1649            "expected split at newline, got: {:?}",
1650            chunks[0]
1651        );
1652        let reassembled: String = chunks.join("");
1653        assert_eq!(reassembled, text);
1654    }
1655
1656    #[test]
1657    fn chunk_text_empty_string() {
1658        let chunks = chunk_text("", 32);
1659        assert!(chunks.is_empty() || chunks == vec!["".to_string()]);
1660    }
1661
1662    // ── build_control_chunks ────────────────────────────────────────────────
1663
1664    #[test]
1665    fn build_control_chunks_session_agent_splits_state_and_message() {
1666        // build_control_chunks always splits state and message into separate chunks
1667        // (it's only called when full_message exceeds the limit)
1668        let state = "state content";
1669        let message = "user question";
1670        let chunks = build_control_chunks(state, message, true);
1671        assert_eq!(chunks.len(), 2);
1672        assert_eq!(chunks[0], "state content");
1673        assert_eq!(chunks[1], "user question");
1674    }
1675
1676    #[test]
1677    fn build_control_chunks_session_agent_long_state_splits() {
1678        // Long state gets split into multiple chunks
1679        let state = "line\n".repeat(10000); // ~60KB
1680        let message = "question";
1681        let chunks = build_control_chunks(&state, message, true);
1682        // Should have multiple state chunks + 1 message chunk
1683        assert!(chunks.len() >= 3);
1684        // All but last are state chunks
1685        for chunk in &chunks[..chunks.len() - 1] {
1686            assert!(chunk.contains("line"));
1687        }
1688        // Last is message
1689        assert_eq!(chunks.last().unwrap(), "question");
1690    }
1691
1692    #[test]
1693    fn build_control_chunks_non_session_agent_adds_headers() {
1694        let state = "state content";
1695        let message = "user question";
1696        let chunks = build_control_chunks(state, message, false);
1697        assert_eq!(chunks.len(), 2);
1698        assert!(chunks[0].starts_with("[Context chunk 1/1]"));
1699        assert!(chunks[0].contains("state content"));
1700        assert!(chunks[1].starts_with("[Context complete. User question:]"));
1701        assert!(chunks[1].contains("user question"));
1702    }
1703
1704    #[test]
1705    fn build_control_chunks_non_session_agent_splits_long_message() {
1706        let state = "state";
1707        let message = "x".repeat(50000); // Exceeds 32KB limit
1708        let chunks = build_control_chunks(state, &message, false);
1709        // Should have 1 state chunk + multiple message chunks
1710        assert!(chunks.len() >= 2);
1711        assert!(chunks[0].starts_with("[Context chunk 1/1]"));
1712        // Message chunks should have question headers
1713        for chunk in &chunks[1..] {
1714            assert!(chunk.starts_with("[User question chunk"));
1715        }
1716        // Verify reassembly (without headers)
1717        let reassembled_state = chunks[0].split_once('\n').map(|(_, r)| r).unwrap_or("");
1718        assert_eq!(reassembled_state, "state");
1719    }
1720
1721    #[test]
1722    fn build_control_chunks_non_session_agent_long_state_splits() {
1723        let state = "line\n".repeat(10000); // ~60KB
1724        let message = "question";
1725        let chunks = build_control_chunks(&state, message, false);
1726        // Multiple state chunks + 1 message chunk
1727        assert!(chunks.len() >= 3);
1728        for (i, chunk) in chunks.iter().enumerate() {
1729            if i < chunks.len() - 1 {
1730                assert!(chunk.starts_with(&format!("[Context chunk {}/", i + 1)));
1731                assert!(chunk.contains("line"));
1732            } else {
1733                assert!(chunk.starts_with("[Context complete. User question:]"));
1734                assert_eq!(
1735                    chunk.split_once('\n').map(|(_, r)| r).unwrap_or(""),
1736                    "question"
1737                );
1738            }
1739        }
1740    }
1741}