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