1use 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
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 '{agent}' not found in PATH. Is it installed?");
169 }
170 Ok(())
171}
172
173#[expect(clippy::print_stderr, reason = "progress line for `orch chat /model`")]
178pub async fn validate_model(spec: &ModelSpec) -> Result<()> {
179 validate_agent(&spec.agent)?;
181
182 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 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
238#[derive(Debug)]
240pub struct ChatContext {
241 pub system: String,
243 pub state: String,
245}
246
247pub async fn assemble_context(store: &TaskStore, session_id: &str) -> Result<ChatContext> {
249 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 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 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 let system = SYSTEM_TEMPLATE.to_string();
332
333 Ok(ChatContext { system, state })
334}
335
336pub 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
346pub 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
356pub 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
364pub 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
373pub 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
382pub 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
389pub 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(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(store, session_id).await?;
438 Ok(format_switched_spec(&spec))
439 }
440 ControlCommand::Restart => {
441 if let Some(ct) = channel_thread {
443 store.kv_set("control:restart_pending", ct).await.ok();
444 }
445 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
459pub 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#[derive(Debug, Clone, PartialEq, Eq)]
476pub struct MemoryEntry {
477 pub key: String,
478 pub value: String,
479}
480
481#[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
496pub 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#[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
566async 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
576pub 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
607fn session_uuid_key(session_name: &str) -> String {
609 format!("control:session_id:{session_name}")
610}
611
612pub 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
641pub 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
656pub 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
694pub 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 if let Some(response) =
707 maybe_handle_control_command(store, session_id, channel_thread, message).await?
708 {
709 return Ok(response);
710 }
711
712 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
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 let ctx = assemble_context(store, session_id).await?;
745
746 let model = get_model(store).await?;
748 let agent = get_agent(store).await?;
749
750 let full_message = format!("{}\n\n---\n\n{}", ctx.state, message);
752
753 let (session_uuid, is_new) = get_or_create_session_uuid(store, session_id).await?;
755
756 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 (
774 build_control_chunks(&ctx.state, message, true),
775 DeliveryMode::Chunked,
776 )
777 } else {
778 (
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 let result = match delivery_mode {
792 DeliveryMode::Chunked => {
793 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 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 let final_parsed = parse_response(&last_result.expect("at least one chunk").text);
902 final_parsed.clean_text
903 }
904 DeliveryMode::ChunkedNonSession => {
905 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 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 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 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
1058 .lock()
1059 .unwrap_or_else(std::sync::PoisonError::into_inner);
1060 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
1073fn 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 while end > start && !text.is_char_boundary(end) {
1087 end -= 1;
1088 }
1089 if end == start {
1090 if let Some(ch) = text[start..].chars().next() {
1092 end = start + ch.len_utf8();
1093 } else {
1094 break;
1095 }
1096 }
1097 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
1109fn 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 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 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 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 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 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 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_agent(&store, "claude").await.unwrap();
1450
1451 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 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 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 #[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 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 #[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 let content = format!("{}😀", "x".repeat(97)); 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 let text = "aaaa\nbbbb";
1648 let chunks = chunk_text(text, 6);
1649 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 #[test]
1668 fn build_control_chunks_session_agent_splits_state_and_message() {
1669 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 let state = "line\n".repeat(10000); let message = "question";
1684 let chunks = build_control_chunks(&state, message, true);
1685 assert!(chunks.len() >= 3);
1687 for chunk in &chunks[..chunks.len() - 1] {
1689 assert!(chunk.contains("line"));
1690 }
1691 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); let chunks = build_control_chunks(state, &message, false);
1712 assert!(chunks.len() >= 2);
1714 assert!(chunks[0].starts_with("[Context chunk 1/1]"));
1715 for chunk in &chunks[1..] {
1717 assert!(chunk.starts_with("[User question chunk"));
1718 }
1719 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); let message = "question";
1728 let chunks = build_control_chunks(&state, message, false);
1729 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}