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