1use 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
37static SESSION_LOCKS: std::sync::LazyLock<Mutex<HashMap<String, Arc<tokio::sync::Mutex<()>>>>> =
47 std::sync::LazyLock::new(|| Mutex::new(HashMap::new()));
48
49const SYSTEM_TEMPLATE: &str = include_str!("../prompts/control_system.md");
51
52const AGENT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(120);
54
55const MAX_CONTROL_MESSAGE_BYTES: usize = 32 * 1024;
59
60enum DeliveryMode {
62 Single,
64 Chunked,
66 ChunkedNonSession,
68}
69
70#[derive(Debug, Clone, PartialEq)]
72pub struct ModelSpec {
73 pub agent: String,
74 pub model: String,
75}
76
77#[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 "grok" => "",
91 _ => "sonnet",
92 }
93}
94
95pub 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
115pub 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
134fn 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"
151 }
152}
153
154pub 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
173pub async fn validate_model(spec: &ModelSpec) -> Result<()> {
178 validate_agent(&spec.agent)?;
180
181 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 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
218async 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
230async fn test_invoke(agent: &str, model: &str) -> Result<()> {
232 let msg = "Reply with just the word 'ok'.";
233 invoke_agent(agent, model, "", msg).await?;
235 Ok(())
236}
237
238pub struct ChatContext {
240 pub system: String,
242 pub state: String,
244}
245
246pub async fn assemble_context(store: &TaskStore, session_id: &str) -> Result<ChatContext> {
248 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 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 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 let system = SYSTEM_TEMPLATE.to_string();
331
332 Ok(ChatContext { system, state })
333}
334
335pub 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
345pub 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
355pub 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
363pub 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
372pub 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
381pub 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
388pub 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(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(store, session_id).await?;
440 Ok(format_switched_spec(&spec))
441 }
442 ControlCommand::Restart => {
443 if let Some(ct) = channel_thread {
445 store.kv_set("control:restart_pending", ct).await.ok();
446 }
447 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
461pub 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#[derive(Debug, Clone, PartialEq, Eq)]
478pub struct MemoryEntry {
479 pub key: String,
480 pub value: String,
481}
482
483#[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
498pub 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
539pub 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
567async 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
577pub 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
608fn session_uuid_key(session_name: &str) -> String {
610 format!("control:session_id:{session_name}")
611}
612
613pub 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
642pub 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
657pub 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
695pub 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 if let Some(response) =
708 maybe_handle_control_command(store, session_id, channel_thread, message).await?
709 {
710 return Ok(response);
711 }
712
713 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
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 let ctx = assemble_context(store, session_id).await?;
744
745 let model = get_model(store).await?;
747 let agent = get_agent(store).await?;
748
749 let full_message = format!("{}\n\n---\n\n{}", ctx.state, message);
751
752 let (session_uuid, is_new) = get_or_create_session_uuid(store, session_id).await?;
754
755 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 (
773 build_control_chunks(&ctx.state, message, true),
774 DeliveryMode::Chunked,
775 )
776 } else {
777 (
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 let result = match delivery_mode {
791 DeliveryMode::Chunked => {
792 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 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 let final_parsed = parse_response(&last_result.expect("at least one chunk").text);
901 final_parsed.clean_text
902 }
903 DeliveryMode::ChunkedNonSession => {
904 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 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 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 let parsed = parse_response(&chunk_result.text);
1002
1003 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 drop(_guard);
1048
1049 if Arc::strong_count(&session_lock) == 2 {
1057 let mut map = SESSION_LOCKS.lock().unwrap_or_else(|e| e.into_inner());
1058 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
1071fn 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 while end > start && !text.is_char_boundary(end) {
1085 end -= 1;
1086 }
1087 if end == start {
1088 if let Some(ch) = text[start..].chars().next() {
1090 end = start + ch.len_utf8();
1091 } else {
1092 break;
1093 }
1094 }
1095 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
1107fn 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 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 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 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 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 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 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_agent(&store, "claude").await.unwrap();
1448
1449 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 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 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 #[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 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 #[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 let content = "x".repeat(97) + "😀"; 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 let text = "aaaa\nbbbb";
1645 let chunks = chunk_text(text, 6);
1646 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 #[test]
1665 fn build_control_chunks_session_agent_splits_state_and_message() {
1666 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 let state = "line\n".repeat(10000); let message = "question";
1681 let chunks = build_control_chunks(&state, message, true);
1682 assert!(chunks.len() >= 3);
1684 for chunk in &chunks[..chunks.len() - 1] {
1686 assert!(chunk.contains("line"));
1687 }
1688 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); let chunks = build_control_chunks(state, &message, false);
1709 assert!(chunks.len() >= 2);
1711 assert!(chunks[0].starts_with("[Context chunk 1/1]"));
1712 for chunk in &chunks[1..] {
1714 assert!(chunk.starts_with("[User question chunk"));
1715 }
1716 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); let message = "question";
1725 let chunks = build_control_chunks(&state, message, false);
1726 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}