1use std::collections::HashSet;
2
3use serde::{Deserialize, Serialize};
4
5use crate::types::{ChatMessage, ImageAttachment, Message, MessageRole, ToolCallMessage};
6
7use crate::types::SessionId;
8
9#[derive(Clone, Debug, Default, Serialize, Deserialize)]
20pub struct RunState {
21 pub turn_tool_calls: usize,
25 pub run_has_tool_calls: bool,
29 pub reasoning_only_strikes: usize,
34 pub empty_response_strikes: usize,
39 pub nudge_count: usize,
43 #[serde(default)]
49 pub thinking_disabled_for_rest_of_run: bool,
50
51 #[serde(default)]
57 pub original_thinking_enabled: bool,
58}
59
60impl RunState {
61 pub fn reset_for_new_run(&mut self) {
63 self.turn_tool_calls = 0;
64 self.run_has_tool_calls = false;
65 self.reasoning_only_strikes = 0;
66 self.empty_response_strikes = 0;
67 self.nudge_count = 0;
68 self.thinking_disabled_for_rest_of_run = false;
69 }
72
73 pub fn record_tool_calls(&mut self, n: usize) {
76 self.turn_tool_calls += n;
77 self.run_has_tool_calls = true;
78 self.reasoning_only_strikes = 0;
79 self.empty_response_strikes = 0;
80 }
81
82 pub fn record_reasoning_only(&mut self) -> usize {
87 self.empty_response_strikes = 0;
88 self.reasoning_only_strikes += 1;
89 if self.reasoning_only_strikes >= 3 {
90 tracing::info!(
91 strikes = self.reasoning_only_strikes,
92 "reasoning_only_strikes reached 3, disabling thinking for rest of run"
93 );
94 self.thinking_disabled_for_rest_of_run = true;
95 }
96 self.reasoning_only_strikes
97 }
98
99 pub fn record_empty_response(&mut self) -> usize {
103 self.reasoning_only_strikes = 0;
104 self.empty_response_strikes += 1;
105 self.empty_response_strikes
106 }
107}
108
109#[derive(Deserialize)]
116struct RawAgentSession {
117 id: Option<SessionId>,
118 chat_messages: Vec<ChatMessage>,
119 always_allowed_actions: HashSet<String>,
120 total_tool_calls: usize,
121
122 run_state: Option<RunState>,
124
125 nudge_count: Option<usize>,
127 turn_tool_calls: Option<usize>,
128 reasoning_only_strikes: Option<usize>,
129 empty_response_strikes: Option<usize>,
130}
131
132impl From<RawAgentSession> for AgentSession {
133 fn from(raw: RawAgentSession) -> Self {
134 let run_state = raw.run_state.unwrap_or_else(|| RunState {
135 nudge_count: raw.nudge_count.unwrap_or(0),
136 turn_tool_calls: raw.turn_tool_calls.unwrap_or(0),
137 reasoning_only_strikes: raw.reasoning_only_strikes.unwrap_or(0),
138 empty_response_strikes: raw.empty_response_strikes.unwrap_or(0),
139 ..RunState::default()
140 });
141 Self {
142 id: raw.id,
143 chat_messages: raw.chat_messages,
144 always_allowed_actions: raw.always_allowed_actions,
145 total_tool_calls: raw.total_tool_calls,
146 run_state,
147 }
148 }
149}
150
151#[derive(Clone, Debug, Default, Serialize, Deserialize)]
156#[serde(from = "RawAgentSession")]
157pub struct AgentSession {
158 id: Option<SessionId>,
159 chat_messages: Vec<ChatMessage>,
162 always_allowed_actions: HashSet<String>,
163 pub total_tool_calls: usize,
166 pub run_state: RunState,
168}
169
170impl AgentSession {
171 pub fn new(id: SessionId) -> Self {
172 Self {
173 id: Some(id),
174 chat_messages: Vec::new(),
175 always_allowed_actions: HashSet::new(),
176 total_tool_calls: 0,
177 run_state: RunState::default(),
178 }
179 }
180
181 pub fn id(&self) -> Option<SessionId> {
182 self.id.clone()
183 }
184
185 pub fn simple_messages(&self) -> Vec<Message> {
189 self.chat_messages
190 .iter()
191 .filter_map(|cm| match cm {
192 ChatMessage::Assistant { content: None, .. } => None,
193 ChatMessage::Assistant {
194 content: Some(c),
195 tool_calls: Some(tc),
196 ..
197 } if c.is_empty() && !tc.is_empty() => None,
198 _ => Some(Message::from(cm)),
199 })
200 .collect()
201 }
202
203 pub fn chat_messages(&self) -> &[ChatMessage] {
204 &self.chat_messages
205 }
206
207 pub fn chat_messages_mut(&mut self) -> &mut Vec<ChatMessage> {
209 &mut self.chat_messages
210 }
211
212 pub fn is_action_allowed(&self, action_key: &str) -> bool {
213 self.always_allowed_actions.contains(action_key)
214 }
215
216 pub fn allow_action(&mut self, action_key: impl Into<String>) {
217 self.always_allowed_actions.insert(action_key.into());
218 }
219
220 pub fn push_message(&mut self, role: MessageRole, content: impl Into<String>) {
221 let content = content.into();
222 let chat_msg = match role {
223 MessageRole::System => ChatMessage::system(content),
224 MessageRole::User => ChatMessage::user(content),
225 MessageRole::Assistant => ChatMessage::assistant(content),
226 MessageRole::Tool => ChatMessage::tool(String::new(), content),
227 };
228 self.chat_messages.push(chat_msg);
229 }
230
231 pub fn push_assistant_with_reasoning(
235 &mut self,
236 content: impl Into<String>,
237 reasoning: impl Into<String>,
238 ) {
239 self.chat_messages
240 .push(ChatMessage::assistant_with_reasoning(content, reasoning));
241 }
242
243 pub fn push_user_message_with_images(
244 &mut self,
245 content: impl Into<String>,
246 images: Vec<ImageAttachment>,
247 ) {
248 self.chat_messages
249 .push(ChatMessage::user_with_images(content, images));
250 }
251
252 pub fn push_assistant_tool_call(
253 &mut self,
254 tool_call_id: &str,
255 tool_name: &str,
256 arguments_json: &str,
257 ) {
258 self.chat_messages.push(ChatMessage::assistant_tool_call(
259 tool_call_id,
260 tool_name,
261 arguments_json,
262 ));
263 }
264
265 pub fn push_assistant_tool_calls(
266 &mut self,
267 tool_calls: &[(String, String, String)],
268 reasoning: Option<String>,
269 content: Option<String>,
270 ) {
271 let calls: Vec<ToolCallMessage> = tool_calls
272 .iter()
273 .map(|(id, name, args)| {
274 let valid_args = if serde_json::from_str::<serde_json::Value>(args).is_ok() {
278 args.clone()
279 } else {
280 tracing::warn!(
281 tool_name = %name,
282 args_len = args.len(),
283 "tool call arguments are not valid JSON (possibly truncated), wrapping in error object"
284 );
285 let max_preview = 200;
288 let safe_end = if args.len() <= max_preview {
289 args.len()
290 } else {
291 args.char_indices()
292 .find(|(i, _)| *i >= max_preview)
293 .map(|(i, _)| i)
294 .unwrap_or(args.len())
295 };
296 serde_json::json!({
297 "error": "tool_call_arguments_truncated",
298 "original_args_preview": &args[..safe_end],
299 "message": "The tool call arguments were truncated or invalid. Please retry with complete arguments."
300 })
301 .to_string()
302 };
303 ToolCallMessage {
304 id: id.clone(),
305 name: name.clone(),
306 arguments: valid_args,
307 }
308 })
309 .collect();
310 self.chat_messages.push(ChatMessage::Assistant {
311 content,
312 reasoning_content: reasoning,
313 tool_calls: Some(calls),
314 thinking_signature: None,
315 });
316 }
317
318 pub fn push_tool_result(&mut self, tool_call_id: &str, content: impl Into<String>) {
319 self.chat_messages
320 .push(ChatMessage::tool(tool_call_id, content));
321 }
322
323 pub fn remove_ephemeral_messages(&mut self) {
327 let before = self.chat_messages.len();
328 self.chat_messages.retain(|m| !m.is_ephemeral());
329 let removed = before - self.chat_messages.len();
330 if removed > 0 {
331 tracing::debug!(
332 removed,
333 remaining = self.chat_messages.len(),
334 "ephemeral messages cleaned up"
335 );
336 }
337 }
338
339 pub fn turn_count(&self) -> usize {
342 self.chat_messages
343 .iter()
344 .filter(|m| matches!(m, ChatMessage::User { .. }))
345 .count()
346 }
347
348 pub fn trim_oldest_turns(&mut self, max_turns: usize) {
351 let current_turns = self.turn_count();
352 if current_turns <= max_turns {
353 return;
354 }
355 let turns_to_remove = current_turns - max_turns;
356
357 let user_positions: Vec<usize> = self
359 .chat_messages
360 .iter()
361 .enumerate()
362 .filter_map(|(i, m)| {
363 if matches!(m, ChatMessage::User { .. }) {
364 Some(i)
365 } else {
366 None
367 }
368 })
369 .collect();
370
371 if user_positions.len() <= turns_to_remove {
372 return;
373 }
374
375 let system_prefix = self
377 .chat_messages
378 .iter()
379 .take_while(|m| matches!(m, ChatMessage::System { .. }))
380 .count();
381
382 let drain_end = user_positions[turns_to_remove];
384 if system_prefix >= drain_end {
385 return; }
387
388 self.chat_messages.drain(system_prefix..drain_end);
389 }
390
391 pub fn pop_last_message(&mut self) {
394 self.chat_messages.pop();
395 }
396
397 pub fn close_dangling_tool_calls(&mut self, error_summary: &str) {
398 let assistant_idx = self.chat_messages.iter().rposition(
399 |m| matches!(m, ChatMessage::Assistant { tool_calls: Some(tc), .. } if !tc.is_empty()),
400 );
401
402 let Some(assistant_idx) = assistant_idx else {
403 return;
404 };
405
406 let ChatMessage::Assistant {
407 tool_calls: Some(tc),
408 ..
409 } = &self.chat_messages[assistant_idx]
410 else {
411 return;
412 };
413
414 let all_ids: Vec<String> = tc.iter().map(|t| t.id.clone()).collect();
415
416 let answered_ids: Vec<String> = self.chat_messages[assistant_idx + 1..]
417 .iter()
418 .filter_map(|m| match m {
419 ChatMessage::Tool { tool_call_id, .. } => Some(tool_call_id.clone()),
420 _ => None,
421 })
422 .collect();
423
424 for id in &all_ids {
425 if !answered_ids.iter().any(|a| a == id) {
426 self.push_tool_result(id, error_summary);
427 }
428 }
429 }
430
431 pub fn set_chat_messages(&mut self, messages: Vec<ChatMessage>) -> Result<(), String> {
436 validate_message_sequence(&messages)?;
437 self.total_tool_calls = messages
440 .iter()
441 .filter_map(|m| match m {
442 ChatMessage::Assistant {
443 tool_calls: Some(tc),
444 ..
445 } => Some(tc.len()),
446 _ => None,
447 })
448 .sum();
449 self.chat_messages = messages;
450 Ok(())
451 }
452}
453
454pub fn validate_message_sequence(messages: &[ChatMessage]) -> Result<(), String> {
462 let mut pending_tool_call_ids: HashSet<String> = HashSet::new();
463
464 for (i, msg) in messages.iter().enumerate() {
465 match msg {
466 ChatMessage::Tool { tool_call_id, .. } => {
467 if pending_tool_call_ids.is_empty() {
468 return Err(format!(
469 "message[{}]: Tool message with call_id '{}' has no preceding tool_call",
470 i, tool_call_id
471 ));
472 }
473 if !pending_tool_call_ids.remove(tool_call_id) {
475 return Err(format!(
476 "message[{}]: Tool message with call_id '{}' does not match any pending tool_call (already answered or unknown)",
477 i, tool_call_id
478 ));
479 }
480 }
481 ChatMessage::Assistant {
482 tool_calls: Some(tc),
483 ..
484 } => {
485 if !pending_tool_call_ids.is_empty() {
487 return Err(format!(
488 "message[{}]: Assistant message with new tool_calls appears before pending calls were answered: {:?}",
489 i, pending_tool_call_ids
490 ));
491 }
492 pending_tool_call_ids = tc.iter().map(|t| t.id.clone()).collect();
493 }
494 _ => {}
495 }
496 }
497
498 if !pending_tool_call_ids.is_empty() {
500 return Err(format!(
501 "message sequence ends with unanswered tool calls: {:?}",
502 pending_tool_call_ids
503 ));
504 }
505
506 Ok(())
507}
508
509#[cfg(test)]
510fn make_session() -> AgentSession {
511 AgentSession::new(SessionId::new(1))
512}
513
514#[cfg(test)]
515mod tests {
516 use super::*;
517
518 #[test]
519 fn test_turn_count_empty() {
520 let s = make_session();
521 assert_eq!(s.turn_count(), 0);
522 }
523
524 #[test]
525 fn test_turn_count_with_system_and_user() {
526 let mut s = make_session();
527 s.push_message(MessageRole::System, "system");
528 assert_eq!(s.turn_count(), 0);
529 s.push_message(MessageRole::User, "hello");
530 assert_eq!(s.turn_count(), 1);
531 s.push_message(MessageRole::Assistant, "hi");
532 assert_eq!(s.turn_count(), 1);
533 s.push_message(MessageRole::User, "bye");
534 assert_eq!(s.turn_count(), 2);
535 }
536
537 #[test]
538 fn test_turn_count_with_tool_calls() {
539 let mut s = make_session();
540 s.push_message(MessageRole::User, "do something");
541 s.push_assistant_tool_calls(&[("id1".into(), "tool".into(), "{}".into())], None, None);
542 s.push_tool_result("id1", "result");
543 s.push_message(MessageRole::Assistant, "done");
544 assert_eq!(s.turn_count(), 1);
546 }
547
548 #[test]
549 fn test_trim_oldest_turns_noop() {
550 let mut s = make_session();
551 s.push_message(MessageRole::User, "hello");
552 s.push_message(MessageRole::Assistant, "hi");
553 s.trim_oldest_turns(5);
554 assert_eq!(s.turn_count(), 1);
555 assert_eq!(s.chat_messages().len(), 2);
556 }
557
558 #[test]
559 fn test_trim_oldest_turns_removes_old() {
560 let mut s = make_session();
561 s.push_message(MessageRole::System, "sys");
562 s.push_message(MessageRole::User, "u1");
564 s.push_message(MessageRole::Assistant, "a1");
565 s.push_message(MessageRole::User, "u2");
567 s.push_message(MessageRole::Assistant, "a2");
568 s.push_message(MessageRole::User, "u3");
570 s.push_message(MessageRole::Assistant, "a3");
571
572 s.trim_oldest_turns(2);
573 assert_eq!(s.turn_count(), 2);
574 assert!(matches!(s.chat_messages()[0], ChatMessage::System { .. }));
576 assert!(
578 matches!(s.chat_messages()[1], ChatMessage::User { ref content, .. } if content == "u2")
579 );
580 }
581
582 #[test]
583 fn test_trim_oldest_turns_with_tool_calls() {
584 let mut s = make_session();
585 s.push_message(MessageRole::User, "u1");
587 s.push_assistant_tool_calls(&[("id1".into(), "t".into(), "{}".into())], None, None);
588 s.push_tool_result("id1", "r1");
589 s.push_message(MessageRole::Assistant, "a1");
590 s.push_message(MessageRole::User, "u2");
592 s.push_message(MessageRole::Assistant, "a2");
593
594 let msg_before = s.simple_messages().len();
595 let chat_before = s.chat_messages().len();
596 s.trim_oldest_turns(1);
597 assert_eq!(s.turn_count(), 1);
598 assert_eq!(s.chat_messages().len(), chat_before - 4);
600 assert_eq!(s.simple_messages().len(), msg_before - 3);
602 }
603
604 #[test]
605 fn test_pop_last_message_text() {
606 let mut s = make_session();
607 s.push_message(MessageRole::User, "hello");
608 s.push_message(MessageRole::Assistant, "hi");
609 assert_eq!(s.chat_messages().len(), 2);
610 s.pop_last_message();
611 assert_eq!(s.chat_messages().len(), 1);
612 assert_eq!(s.simple_messages().len(), 1);
613 }
614
615 #[test]
616 fn test_pop_last_message_tool_calls_only() {
617 let mut s = make_session();
618 s.push_message(MessageRole::User, "do it");
619 s.push_assistant_tool_calls(&[("id1".into(), "t".into(), "{}".into())], None, None);
620 assert_eq!(s.chat_messages().len(), 2);
621 assert_eq!(s.simple_messages().len(), 1); s.pop_last_message();
623 assert_eq!(s.chat_messages().len(), 1);
624 assert_eq!(s.simple_messages().len(), 1); }
626
627 #[test]
628 fn test_pop_last_message_empty_session() {
629 let mut s = make_session();
630 s.pop_last_message(); assert_eq!(s.chat_messages().len(), 0);
632 }
633
634 #[test]
637 fn test_id_and_action_allowlist() {
638 let mut s = make_session();
639 assert_eq!(s.id(), Some(SessionId::new(1)));
640 assert!(!s.is_action_allowed("approve:rm"));
641 s.allow_action("approve:rm");
642 assert!(s.is_action_allowed("approve:rm"));
643 assert!(!s.is_action_allowed("approve:shell"));
644 }
645
646 #[test]
647 fn test_chat_messages_mut() {
648 let mut s = make_session();
649 s.chat_messages_mut().push(ChatMessage::user("direct"));
650 assert_eq!(s.chat_messages().len(), 1);
651 }
652
653 #[test]
654 fn test_push_message_tool_role() {
655 let mut s = make_session();
656 s.push_message(MessageRole::Tool, "result");
657 assert!(matches!(s.chat_messages()[0], ChatMessage::Tool { .. }));
658 }
659
660 #[test]
661 fn test_push_assistant_with_reasoning() {
662 let mut s = make_session();
663 s.push_assistant_with_reasoning("answer", "thinking");
664 match &s.chat_messages()[0] {
665 ChatMessage::Assistant {
666 content,
667 reasoning_content,
668 ..
669 } => {
670 assert_eq!(content.as_deref(), Some("answer"));
671 assert_eq!(reasoning_content.as_deref(), Some("thinking"));
672 }
673 other => panic!("unexpected message: {other:?}"),
674 }
675 }
676
677 #[test]
678 fn test_push_user_message_with_images() {
679 let mut s = make_session();
680 s.push_user_message_with_images(
681 "look",
682 vec![ImageAttachment::Url {
683 url: "http://x".into(),
684 detail: None,
685 }],
686 );
687 match &s.chat_messages()[0] {
688 ChatMessage::User { images, .. } => assert_eq!(images.len(), 1),
689 other => panic!("unexpected message: {other:?}"),
690 }
691 }
692
693 #[test]
694 fn test_push_assistant_tool_call_singular() {
695 let mut s = make_session();
696 s.push_assistant_tool_call("call_1", "bash", "{}");
697 match &s.chat_messages()[0] {
698 ChatMessage::Assistant {
699 tool_calls: Some(tc),
700 ..
701 } => {
702 assert_eq!(tc.len(), 1);
703 assert_eq!(tc[0].id, "call_1");
704 assert_eq!(tc[0].name, "bash");
705 }
706 other => panic!("unexpected message: {other:?}"),
707 }
708 }
709
710 #[test]
711 fn test_simple_messages_filters_empty_content_tool_calls() {
712 let mut s = make_session();
713 s.chat_messages_mut().push(ChatMessage::Assistant {
714 content: Some(String::new()),
715 reasoning_content: None,
716 tool_calls: Some(vec![ToolCallMessage {
717 id: "c".into(),
718 name: "t".into(),
719 arguments: "{}".into(),
720 }]),
721 thinking_signature: None,
722 });
723 assert!(s.simple_messages().is_empty());
724 }
725
726 #[test]
727 fn test_remove_ephemeral_messages() {
728 let mut s = make_session();
729 s.push_message(MessageRole::System, "keep");
730 s.chat_messages_mut()
731 .push(ChatMessage::user_ephemeral("temp"));
732 s.chat_messages_mut()
733 .push(ChatMessage::system_ephemeral("temp2"));
734 s.push_message(MessageRole::User, "keep2");
735 assert_eq!(s.chat_messages().len(), 4);
736 s.remove_ephemeral_messages();
737 assert_eq!(s.chat_messages().len(), 2);
738 assert!(s.chat_messages().iter().all(|m| !m.is_ephemeral()));
739 }
740
741 #[test]
742 fn test_close_dangling_tool_calls_noop_without_tool_call() {
743 let mut s = make_session();
744 s.push_message(MessageRole::User, "hi");
745 s.push_message(MessageRole::Assistant, "hi");
746 s.close_dangling_tool_calls("failed");
747 assert_eq!(s.chat_messages().len(), 2);
748 }
749
750 #[test]
751 fn test_close_dangling_tool_calls_adds_missing_results() {
752 let mut s = make_session();
753 s.push_message(MessageRole::User, "do");
754 s.push_assistant_tool_calls(
755 &[
756 ("c1".into(), "t".into(), "{}".into()),
757 ("c2".into(), "t".into(), "{}".into()),
758 ],
759 None,
760 None,
761 );
762 s.push_tool_result("c1", "ok"); s.close_dangling_tool_calls("failed");
764
765 let tool_results: Vec<(String, String)> = s
766 .chat_messages()
767 .iter()
768 .filter_map(|m| match m {
769 ChatMessage::Tool {
770 tool_call_id,
771 name: _,
772 content,
773 } => Some((tool_call_id.clone(), content.clone())),
774 _ => None,
775 })
776 .collect();
777 assert_eq!(tool_results.len(), 2);
778 assert!(
779 tool_results
780 .iter()
781 .any(|(id, c)| id == "c2" && c == "failed")
782 );
783 }
784
785 #[test]
786 fn test_set_chat_messages_recalculates_total_tool_calls() {
787 let mut s = make_session();
788 let msgs = vec![
789 ChatMessage::user("do"),
790 ChatMessage::assistant_tool_call("c1", "t", "{}"),
791 ChatMessage::tool("c1", "result"),
792 ];
793 s.set_chat_messages(msgs).unwrap();
794 assert_eq!(s.total_tool_calls, 1);
795 }
796}
797
798#[cfg(test)]
799mod validate_tests {
800 use super::*;
801
802 #[test]
803 fn test_valid_simple_sequence() {
804 let msgs = vec![ChatMessage::user("hello"), ChatMessage::assistant("hi")];
805 assert!(validate_message_sequence(&msgs).is_ok());
806 }
807
808 #[test]
809 fn test_valid_tool_call_sequence() {
810 let msgs = vec![
811 ChatMessage::user("run command"),
812 ChatMessage::assistant_tool_call("call_1", "bash", r#"{"cmd":"ls"}"#),
813 ChatMessage::tool("call_1", "file1 file2"),
814 ChatMessage::assistant("done"),
815 ];
816 assert!(validate_message_sequence(&msgs).is_ok());
817 }
818
819 #[test]
820 fn test_valid_multi_tool_call_sequence() {
821 let msgs = vec![
822 ChatMessage::user("run commands"),
823 ChatMessage::Assistant {
824 content: None,
825 reasoning_content: None,
826 tool_calls: Some(vec![
827 crate::types::ToolCallMessage {
828 id: "call_1".into(),
829 name: "bash".into(),
830 arguments: "{}".into(),
831 },
832 crate::types::ToolCallMessage {
833 id: "call_2".into(),
834 name: "read".into(),
835 arguments: "{}".into(),
836 },
837 ]),
838 thinking_signature: None,
839 },
840 ChatMessage::tool("call_1", "result1"),
841 ChatMessage::tool("call_2", "result2"),
842 ChatMessage::assistant("done"),
843 ];
844 assert!(validate_message_sequence(&msgs).is_ok());
845 }
846
847 #[test]
848 fn test_orphaned_tool_result() {
849 let msgs = vec![
850 ChatMessage::user("hello"),
851 ChatMessage::tool("call_1", "orphaned result"),
852 ];
853 let err = validate_message_sequence(&msgs).unwrap_err();
854 assert!(err.contains("no preceding tool_call"));
855 }
856
857 #[test]
858 fn test_mismatched_tool_call_id() {
859 let msgs = vec![
860 ChatMessage::user("run"),
861 ChatMessage::assistant_tool_call("call_1", "bash", "{}"),
862 ChatMessage::tool("call_2", "wrong id"),
863 ];
864 let err = validate_message_sequence(&msgs).unwrap_err();
865 assert!(err.contains("does not match"));
866 }
867
868 #[test]
869 fn test_set_chat_messages_valid() {
870 let mut s = make_session();
871 let msgs = vec![ChatMessage::user("hello"), ChatMessage::assistant("hi")];
872 assert!(s.set_chat_messages(msgs.clone()).is_ok());
873 assert_eq!(s.chat_messages().len(), 2);
874 }
875
876 #[test]
877 fn test_set_chat_messages_invalid() {
878 let mut s = make_session();
879 let msgs = vec![ChatMessage::tool("call_1", "orphaned")];
880 assert!(s.set_chat_messages(msgs).is_err());
881 }
882
883 #[test]
886 fn run_state_default() {
887 let rs = RunState::default();
888 assert_eq!(rs.turn_tool_calls, 0);
889 assert!(!rs.run_has_tool_calls);
890 assert_eq!(rs.reasoning_only_strikes, 0);
891 assert_eq!(rs.empty_response_strikes, 0);
892 assert_eq!(rs.nudge_count, 0);
893 }
894
895 #[test]
896 fn run_state_reset_for_new_run() {
897 let mut rs = RunState {
898 turn_tool_calls: 5,
899 run_has_tool_calls: true,
900 reasoning_only_strikes: 2,
901 empty_response_strikes: 1,
902 nudge_count: 3,
903 thinking_disabled_for_rest_of_run: true,
904 original_thinking_enabled: true,
905 };
906
907 rs.reset_for_new_run();
908
909 assert_eq!(rs.turn_tool_calls, 0);
910 assert!(!rs.run_has_tool_calls);
911 assert_eq!(rs.reasoning_only_strikes, 0);
912 assert_eq!(rs.empty_response_strikes, 0);
913 assert_eq!(rs.nudge_count, 0);
914 assert!(!rs.thinking_disabled_for_rest_of_run);
915 assert!(rs.original_thinking_enabled);
917 }
918
919 #[test]
920 fn run_state_record_tool_calls() {
921 let mut rs = RunState {
922 reasoning_only_strikes: 2,
923 empty_response_strikes: 1,
924 ..RunState::default()
925 };
926
927 rs.record_tool_calls(3);
928
929 assert_eq!(rs.turn_tool_calls, 3);
930 assert!(rs.run_has_tool_calls);
931 assert_eq!(rs.reasoning_only_strikes, 0); assert_eq!(rs.empty_response_strikes, 0); }
934
935 #[test]
936 fn run_state_record_tool_calls_accumulates() {
937 let mut rs = RunState::default();
938 rs.record_tool_calls(2);
939 rs.record_tool_calls(3);
940
941 assert_eq!(rs.turn_tool_calls, 5);
942 assert!(rs.run_has_tool_calls);
943 }
944
945 #[test]
946 fn run_state_record_reasoning_only() {
947 let mut rs = RunState {
948 empty_response_strikes: 2,
949 ..RunState::default()
950 };
951
952 let strikes = rs.record_reasoning_only();
953
954 assert_eq!(strikes, 1);
955 assert_eq!(rs.reasoning_only_strikes, 1);
956 assert_eq!(rs.empty_response_strikes, 0); }
958
959 #[test]
960 fn run_state_record_reasoning_only_consecutive() {
961 let mut rs = RunState::default();
962
963 assert_eq!(rs.record_reasoning_only(), 1);
964 assert_eq!(rs.record_reasoning_only(), 2);
965 assert_eq!(rs.record_reasoning_only(), 3);
966 }
967
968 #[test]
969 fn run_state_record_empty_response() {
970 let mut rs = RunState {
971 reasoning_only_strikes: 2,
972 ..RunState::default()
973 };
974
975 let strikes = rs.record_empty_response();
976
977 assert_eq!(strikes, 1);
978 assert_eq!(rs.empty_response_strikes, 1);
979 assert_eq!(rs.reasoning_only_strikes, 0); }
981
982 #[test]
983 fn run_state_record_empty_response_consecutive() {
984 let mut rs = RunState::default();
985
986 assert_eq!(rs.record_empty_response(), 1);
987 assert_eq!(rs.record_empty_response(), 2);
988 assert_eq!(rs.record_empty_response(), 3);
989 }
990
991 #[test]
992 fn run_state_branch_cross_reset() {
993 let mut rs = RunState::default();
995
996 rs.record_reasoning_only();
998 assert_eq!(rs.reasoning_only_strikes, 1);
999
1000 rs.record_tool_calls(2);
1002 assert_eq!(rs.reasoning_only_strikes, 0);
1003 assert_eq!(rs.turn_tool_calls, 2);
1004
1005 let strikes = rs.record_reasoning_only();
1007 assert_eq!(strikes, 1);
1008 }
1009
1010 #[test]
1011 fn run_state_empty_to_reasoning_reset() {
1012 let mut rs = RunState::default();
1014
1015 rs.record_empty_response();
1016 rs.record_empty_response();
1017 assert_eq!(rs.empty_response_strikes, 2);
1018
1019 rs.record_reasoning_only();
1021 assert_eq!(rs.empty_response_strikes, 0);
1022 assert_eq!(rs.reasoning_only_strikes, 1);
1023 }
1024
1025 #[test]
1026 fn run_state_thinking_disabled_default() {
1027 let rs = RunState::default();
1028 assert!(!rs.thinking_disabled_for_rest_of_run);
1029 }
1030
1031 #[test]
1032 fn run_state_thinking_disabled_after_3_strikes() {
1033 let mut rs = RunState::default();
1034
1035 rs.record_reasoning_only();
1037 assert!(!rs.thinking_disabled_for_rest_of_run);
1038 assert_eq!(rs.reasoning_only_strikes, 1);
1039
1040 rs.record_reasoning_only();
1042 assert!(!rs.thinking_disabled_for_rest_of_run);
1043 assert_eq!(rs.reasoning_only_strikes, 2);
1044
1045 rs.record_reasoning_only();
1047 assert!(rs.thinking_disabled_for_rest_of_run);
1048 assert_eq!(rs.reasoning_only_strikes, 3);
1049 }
1050
1051 #[test]
1052 fn run_state_thinking_disabled_resets_on_new_run() {
1053 let mut rs = RunState::default();
1054
1055 rs.record_reasoning_only();
1057 rs.record_reasoning_only();
1058 rs.record_reasoning_only();
1059 assert!(rs.thinking_disabled_for_rest_of_run);
1060 assert_eq!(rs.reasoning_only_strikes, 3);
1061
1062 rs.reset_for_new_run();
1064 assert!(!rs.thinking_disabled_for_rest_of_run);
1065 assert_eq!(rs.reasoning_only_strikes, 0);
1066 }
1067
1068 #[test]
1069 fn run_state_thinking_disabled_stays_after_tool_calls() {
1070 let mut rs = RunState::default();
1071
1072 rs.record_reasoning_only();
1074 rs.record_reasoning_only();
1075 rs.record_reasoning_only();
1076 assert!(rs.thinking_disabled_for_rest_of_run);
1077
1078 rs.record_tool_calls(2);
1081 assert!(rs.thinking_disabled_for_rest_of_run);
1082 assert_eq!(rs.reasoning_only_strikes, 0); }
1084
1085 #[test]
1088 fn deserialize_legacy_flat_fields() {
1089 let json = r#"{
1091 "id": null,
1092 "chat_messages": [],
1093 "always_allowed_actions": [],
1094 "total_tool_calls": 5,
1095 "nudge_count": 3,
1096 "turn_tool_calls": 2,
1097 "reasoning_only_strikes": 1,
1098 "empty_response_strikes": 0
1099 }"#;
1100 let session: AgentSession = serde_json::from_str(json).unwrap();
1101 assert_eq!(session.run_state.nudge_count, 3);
1102 assert_eq!(session.run_state.turn_tool_calls, 2);
1103 assert_eq!(session.run_state.reasoning_only_strikes, 1);
1104 assert_eq!(session.run_state.empty_response_strikes, 0);
1105 assert!(!session.run_state.run_has_tool_calls); }
1107
1108 #[test]
1109 fn deserialize_new_run_state_format() {
1110 let json = r#"{
1112 "id": null,
1113 "chat_messages": [],
1114 "always_allowed_actions": [],
1115 "total_tool_calls": 5,
1116 "run_state": {
1117 "turn_tool_calls": 4,
1118 "run_has_tool_calls": true,
1119 "reasoning_only_strikes": 0,
1120 "empty_response_strikes": 1,
1121 "nudge_count": 2
1122 }
1123 }"#;
1124 let session: AgentSession = serde_json::from_str(json).unwrap();
1125 assert_eq!(session.run_state.turn_tool_calls, 4);
1126 assert!(session.run_state.run_has_tool_calls);
1127 assert_eq!(session.run_state.empty_response_strikes, 1);
1128 assert_eq!(session.run_state.nudge_count, 2);
1129 }
1130
1131 #[test]
1132 fn deserialize_run_state_takes_precedence_over_flat() {
1133 let json = r#"{
1135 "id": null,
1136 "chat_messages": [],
1137 "always_allowed_actions": [],
1138 "total_tool_calls": 0,
1139 "run_state": {
1140 "turn_tool_calls": 10,
1141 "run_has_tool_calls": true,
1142 "reasoning_only_strikes": 0,
1143 "empty_response_strikes": 0,
1144 "nudge_count": 0
1145 },
1146 "nudge_count": 99,
1147 "turn_tool_calls": 99
1148 }"#;
1149 let session: AgentSession = serde_json::from_str(json).unwrap();
1150 assert_eq!(session.run_state.turn_tool_calls, 10); assert_eq!(session.run_state.nudge_count, 0); }
1153
1154 #[test]
1155 fn deserialize_legacy_missing_optional_fields() {
1156 let json = r#"{
1158 "id": null,
1159 "chat_messages": [],
1160 "always_allowed_actions": [],
1161 "total_tool_calls": 0,
1162 "nudge_count": 1
1163 }"#;
1164 let session: AgentSession = serde_json::from_str(json).unwrap();
1165 assert_eq!(session.run_state.nudge_count, 1);
1166 assert_eq!(session.run_state.turn_tool_calls, 0); assert_eq!(session.run_state.reasoning_only_strikes, 0);
1168 assert_eq!(session.run_state.empty_response_strikes, 0);
1169 }
1170
1171 #[test]
1172 fn roundtrip_preserves_run_state() {
1173 let mut session = AgentSession::new(SessionId::new(1));
1174 session.run_state.nudge_count = 5;
1175 session.run_state.turn_tool_calls = 3;
1176 session.run_state.run_has_tool_calls = true;
1177 session.run_state.reasoning_only_strikes = 2;
1178
1179 let json = serde_json::to_string(&session).unwrap();
1180 let restored: AgentSession = serde_json::from_str(&json).unwrap();
1181 assert_eq!(restored.run_state.nudge_count, 5);
1182 assert_eq!(restored.run_state.turn_tool_calls, 3);
1183 assert!(restored.run_state.run_has_tool_calls);
1184 assert_eq!(restored.run_state.reasoning_only_strikes, 2);
1185 }
1186
1187 #[test]
1188 fn push_assistant_tool_calls_validates_json_args() {
1189 let mut s = make_session();
1190
1191 let valid_args = r#"{"path": "src/main.rs", "content": "fn main() {}"}"#;
1193 s.push_assistant_tool_calls(
1194 &[("id1".into(), "write_file".into(), valid_args.into())],
1195 None,
1196 None,
1197 );
1198 if let ChatMessage::Assistant {
1199 tool_calls: Some(ref tc),
1200 ..
1201 } = s.chat_messages[0]
1202 {
1203 assert_eq!(tc[0].arguments, valid_args);
1204 } else {
1205 panic!("expected Assistant message with tool_calls");
1206 }
1207
1208 let truncated_args = r#"{"path": "src/ui/markdown.rs", "content": "#;
1210 s.push_assistant_tool_calls(
1211 &[("id2".into(), "write_file".into(), truncated_args.into())],
1212 None,
1213 None,
1214 );
1215 if let ChatMessage::Assistant {
1216 tool_calls: Some(ref tc),
1217 ..
1218 } = s.chat_messages[1]
1219 {
1220 let parsed: serde_json::Value = serde_json::from_str(&tc[0].arguments)
1222 .expect("wrapped arguments should be valid JSON");
1223 assert_eq!(parsed["error"], "tool_call_arguments_truncated");
1224 assert!(parsed["message"].as_str().unwrap().contains("truncated"));
1225 } else {
1226 panic!("expected Assistant message with tool_calls");
1227 }
1228 }
1229
1230 #[test]
1231 fn push_assistant_tool_calls_truncated_multibyte_no_panic() {
1232 let mut s = make_session();
1235
1236 let mut bad_args = "あ".repeat(70); bad_args.push_str("truncated"); s.push_assistant_tool_calls(&[("id1".into(), "tool".into(), bad_args)], None, None);
1243
1244 if let ChatMessage::Assistant {
1245 tool_calls: Some(ref tc),
1246 ..
1247 } = s.chat_messages[0]
1248 {
1249 let parsed: serde_json::Value = serde_json::from_str(&tc[0].arguments)
1251 .expect("wrapped arguments should be valid JSON even with multibyte chars");
1252 assert_eq!(parsed["error"], "tool_call_arguments_truncated");
1253 } else {
1254 panic!("expected Assistant message with tool_calls");
1255 }
1256 }
1257
1258 #[test]
1259 fn push_assistant_tool_calls_then_tool_result_matches_anthropic_protocol() {
1260 let mut s = make_session();
1265
1266 let tool_calls = vec![
1268 (
1269 "call_00_VJtlnKha0ZZ2Yo8t5ysQ8883".to_string(),
1270 "write_file".to_string(),
1271 "{}".to_string(),
1272 ),
1273 (
1274 "call_01_abc123".to_string(),
1275 "bash".to_string(),
1276 r#"{"command": "ls"}"#.to_string(),
1277 ),
1278 ];
1279
1280 s.push_assistant_tool_calls(
1282 &tool_calls,
1283 Some("thinking...".to_string()),
1284 Some("I'll help you".to_string()),
1285 );
1286
1287 for (tc_id, _, _) in &tool_calls {
1289 s.push_tool_result(
1290 tc_id,
1291 "Tool call was not executed: the response hit the output token limit.",
1292 );
1293 }
1294
1295 if let ChatMessage::Assistant {
1298 tool_calls: Some(ref tc),
1299 ..
1300 } = s.chat_messages[0]
1301 {
1302 assert_eq!(tc.len(), 2);
1303 assert_eq!(tc[0].id, "call_00_VJtlnKha0ZZ2Yo8t5ysQ8883");
1304 assert_eq!(tc[0].name, "write_file");
1305 assert_eq!(tc[1].id, "call_01_abc123");
1306 assert_eq!(tc[1].name, "bash");
1307 } else {
1308 panic!("expected Assistant message with tool_calls");
1309 }
1310
1311 if let ChatMessage::Tool { tool_call_id, .. } = &s.chat_messages[1] {
1313 assert_eq!(tool_call_id, "call_00_VJtlnKha0ZZ2Yo8t5ysQ8883");
1314 } else {
1315 panic!("expected Tool message for first tool call");
1316 }
1317
1318 if let ChatMessage::Tool { tool_call_id, .. } = &s.chat_messages[2] {
1319 assert_eq!(tool_call_id, "call_01_abc123");
1320 } else {
1321 panic!("expected Tool message for second tool call");
1322 }
1323
1324 assert!(
1326 validate_message_sequence(&s.chat_messages).is_ok(),
1327 "message sequence should be valid with matching tool_use and tool_result"
1328 );
1329 }
1330}
1331
1332#[cfg(test)]
1333mod proptest_tests {
1334 use super::*;
1335 use proptest::prelude::*;
1336
1337 proptest! {
1338 #[test]
1341 fn reset_for_new_run_zeros_all_fields(
1342 turn_tool_calls in 0usize..1000,
1343 run_has_tool_calls in proptest::bool::ANY,
1344 reasoning_only_strikes in 0usize..100,
1345 empty_response_strikes in 0usize..100,
1346 nudge_count in 0usize..100,
1347 ) {
1348 let mut rs = RunState {
1349 turn_tool_calls,
1350 run_has_tool_calls,
1351 reasoning_only_strikes,
1352 empty_response_strikes,
1353 nudge_count,
1354 thinking_disabled_for_rest_of_run: true,
1355 original_thinking_enabled: true,
1356 };
1357 rs.reset_for_new_run();
1358 assert_eq!(rs.turn_tool_calls, 0);
1359 assert!(!rs.run_has_tool_calls);
1360 assert_eq!(rs.reasoning_only_strikes, 0);
1361 assert_eq!(rs.empty_response_strikes, 0);
1362 assert_eq!(rs.nudge_count, 0);
1363 assert!(!rs.thinking_disabled_for_rest_of_run);
1364 assert!(rs.original_thinking_enabled);
1366 }
1367
1368 #[test]
1369 fn record_tool_calls_accumulates(n in 0usize..100) {
1370 let mut rs = RunState::default();
1371 rs.record_tool_calls(n);
1372 assert_eq!(rs.turn_tool_calls, n);
1373 assert!(rs.run_has_tool_calls);
1374 assert_eq!(rs.reasoning_only_strikes, 0);
1375 assert_eq!(rs.empty_response_strikes, 0);
1376 }
1377
1378 #[test]
1379 fn record_reasoning_only_increments(count in 1usize..50) {
1380 let mut rs = RunState::default();
1381 for i in 1..=count {
1382 let strikes = rs.record_reasoning_only();
1383 assert_eq!(strikes, i);
1384 assert_eq!(rs.empty_response_strikes, 0);
1385 }
1386 }
1387
1388 #[test]
1389 fn record_empty_response_increments(count in 1usize..50) {
1390 let mut rs = RunState::default();
1391 for i in 1..=count {
1392 let strikes = rs.record_empty_response();
1393 assert_eq!(strikes, i);
1394 assert_eq!(rs.reasoning_only_strikes, 0);
1395 }
1396 }
1397
1398 #[test]
1401 fn push_assistant_tool_calls_valid_json_unchanged(args in r"\{[^{}]{0,200}\}") {
1402 if serde_json::from_str::<serde_json::Value>(&args).is_err() {
1404 return Ok(());
1405 }
1406 let mut s = make_session();
1407 s.push_assistant_tool_calls(
1408 &[("id".into(), "tool".into(), args.clone())],
1409 None,
1410 None,
1411 );
1412 if let ChatMessage::Assistant { tool_calls: Some(ref tc), .. } = s.chat_messages()[0] {
1413 assert_eq!(tc[0].arguments, args);
1414 } else {
1415 panic!("expected Assistant with tool_calls");
1416 }
1417 }
1418
1419 #[test]
1420 fn push_assistant_tool_calls_invalid_json_wrapped_safely(
1421 bad_args in "[a-z\u{4e00}-\u{9fff}]{0,300}"
1422 ) {
1423 if serde_json::from_str::<serde_json::Value>(&bad_args).is_ok() {
1425 return Ok(());
1426 }
1427 let mut s = make_session();
1428 s.push_assistant_tool_calls(
1429 &[("id".into(), "tool".into(), bad_args)],
1430 None,
1431 None,
1432 );
1433 if let ChatMessage::Assistant { tool_calls: Some(ref tc), .. } = s.chat_messages()[0] {
1434 let parsed: serde_json::Value = serde_json::from_str(&tc[0].arguments)
1435 .expect("wrapped args must be valid JSON");
1436 assert_eq!(parsed["error"], "tool_call_arguments_truncated");
1437 } else {
1438 panic!("expected Assistant with tool_calls");
1439 }
1440 }
1441
1442 #[test]
1445 fn trim_oldest_turns_never_exceeds_max(turns in 1usize..20, max in 1usize..20) {
1446 let mut s = make_session();
1447 for i in 0..turns {
1448 s.push_message(MessageRole::User, format!("u{}", i));
1449 s.push_message(MessageRole::Assistant, format!("a{}", i));
1450 }
1451 s.trim_oldest_turns(max);
1452 assert!(s.turn_count() <= max || turns <= max);
1453 }
1454
1455 #[test]
1458 fn validate_simple_user_assistant_always_passes(count in 1usize..20) {
1459 let mut msgs = Vec::new();
1460 for i in 0..count {
1461 msgs.push(ChatMessage::user(format!("msg{}", i)));
1462 msgs.push(ChatMessage::assistant(format!("reply{}", i)));
1463 }
1464 assert!(validate_message_sequence(&msgs).is_ok());
1465 }
1466 }
1467}