Skip to main content

agent_base/engine/
session.rs

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/// Run-level state tracking for the react loop.
10///
11/// Manages counters and flags that track the current run's progress.
12/// All fields are reset at the start of each run (when a new user message arrives).
13///
14/// # Backward compatibility
15///
16/// Before this struct existed, `nudge_count` and `turn_tool_calls` were flat
17/// fields on `AgentSession`.  The [`RawAgentSession`] deserialization shim
18/// migrates them automatically.
19#[derive(Clone, Debug, Default, Serialize, Deserialize)]
20pub struct RunState {
21    /// Number of tool calls already executed in the current turn.
22    /// Reset to 0 at the start of each turn (when a new user message arrives).
23    /// Used by `TurnToolLimitMiddleware` to enforce per-turn tool call limits.
24    pub turn_tool_calls: usize,
25    /// Whether any tools were called in the current run.
26    /// Reset to false at the start of each run. Used by the completion judge
27    /// to determine if tools were used before a text-only response.
28    pub run_has_tool_calls: bool,
29    /// Number of consecutive LLM turns that produced only reasoning_content
30    /// (no text, no tool call). Reset to 0 when a normal response or tool call
31    /// is produced. Used by the react loop to fail instead of looping forever
32    /// on a reasoning-model runaway.
33    pub reasoning_only_strikes: usize,
34    /// Number of consecutive LLM turns that produced a completely empty response
35    /// (no text, no reasoning, no tool call). Reset to 0 when a normal response
36    /// or tool call is produced. Used by the react loop to retry a bounded
37    /// number of times, then fail instead of looping forever.
38    pub empty_response_strikes: usize,
39    /// Number of tool-enforcement nudges issued in the current turn.
40    /// Reset to 0 at the start of each turn (when a new user message arrives).
41    /// Used by `ToolEnforcementMiddleware` to cap nudge attempts per turn.
42    pub nudge_count: usize,
43    /// Thinking is disabled for the rest of the current run.
44    /// Set to true when reasoning_only_strikes reaches the maximum (3).
45    /// Reset to false at the start of each new run (when a new user message arrives).
46    /// This allows the model to continue without thinking after too many
47    /// reasoning-only responses.
48    #[serde(default)]
49    pub thinking_disabled_for_rest_of_run: bool,
50
51    // ─── New fields for thinking guard ─────────────────────────────
52    /// Original thinking configuration (for restoration)
53    ///
54    /// Records whether thinking was originally enabled when the session started.
55    /// Used to restore thinking to its original state after DisableThinking → RestoreThinking cycle.
56    #[serde(default)]
57    pub original_thinking_enabled: bool,
58}
59
60impl RunState {
61    /// Reset all run-level state for a new run (when a new user message arrives).
62    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        // Note: original_thinking_enabled is NOT reset here
70        // It should be set once when the session starts
71    }
72
73    /// Record tool calls (branch 3: tool calls).
74    /// Resets reasoning_only_strikes and empty_response_strikes.
75    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    /// Record reasoning-only response (branch 1).
83    /// Resets empty_response_strikes.
84    /// Returns the new strike count.
85    /// If strikes reach 3, disables thinking for the rest of the run.
86    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    /// Record empty response (branch 2).
100    /// Resets reasoning_only_strikes.
101    /// Returns the new strike count.
102    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/// Deserialization shim for backward compatibility.
110///
111/// Before `RunState` was introduced, `nudge_count`, `turn_tool_calls`,
112/// `reasoning_only_strikes`, and `empty_response_strikes` were flat fields
113/// on `AgentSession`.  This struct accepts **both** the old flat format and
114/// the new nested `run_state` format, migrating legacy data on the fly.
115#[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    // ── new format (preferred) ──
123    run_state: Option<RunState>,
124
125    // ── legacy flat fields (fallback) ──
126    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/// Stable identity + rolling state of a single chat thread.
152///
153/// `Default` returns a fresh session; load persisted state via serde.
154/// Backward-compatible with the old flat-field format (pre-`RunState` migration).
155#[derive(Clone, Debug, Default, Serialize, Deserialize)]
156#[serde(from = "RawAgentSession")]
157pub struct AgentSession {
158    id: Option<SessionId>,
159    /// LLM API format messages, sent directly to the provider.
160    /// This is the single source of truth for the conversation state.
161    chat_messages: Vec<ChatMessage>,
162    always_allowed_actions: HashSet<String>,
163    /// Total number of tool calls made in this session (across all turns).
164    /// Used by middleware for decisions like "first_turn_only" enforcement.
165    pub total_tool_calls: usize,
166    /// Run-level state tracking for the react loop.
167    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    /// Derive a simplified `Vec<Message>` view from the canonical `chat_messages`.
186    /// Assistant messages that contain only tool_calls (no text content) are
187    /// skipped, since they have no corresponding simplified representation.
188    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    /// 可变引用,仅用于需要直接操作消息的高级场景。
208    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    /// Push an assistant message with reasoning/thinking content preserved.
232    /// This allows the LLM to see its own prior reasoning in subsequent turns,
233    /// preventing it from re-deriving the same conclusions every turn.
234    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                // Validate that arguments are valid JSON.
275                // If not (e.g., truncated by token limit), wrap in an error object
276                // so the API doesn't reject the entire request with 400.
277                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                    // Find a safe char boundary at or before byte 200 to avoid
286                    // panicking on multi-byte UTF-8 chars (CJK, emoji, etc.).
287                    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    /// 移除所有临时消息(ephemeral=true)。
324    ///
325    /// 在 turn 结束时调用,确保注入的临时内容不残留到下一轮。
326    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    /// Count the number of conversation turns.
340    /// A turn starts with a User message and includes subsequent Assistant/Tool messages.
341    pub fn turn_count(&self) -> usize {
342        self.chat_messages
343            .iter()
344            .filter(|m| matches!(m, ChatMessage::User { .. }))
345            .count()
346    }
347
348    /// Remove the oldest turns from the front until turn count ≤ max_turns.
349    /// Preserves the System message at index 0 if present.
350    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        // Find User message positions (turn boundaries) in chat_messages
358        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        // Preserve system prefix: count leading System messages
376        let system_prefix = self
377            .chat_messages
378            .iter()
379            .take_while(|m| matches!(m, ChatMessage::System { .. }))
380            .count();
381
382        // Drain from system_prefix up to the start of the (turns_to_remove + 1)-th turn
383        let drain_end = user_positions[turns_to_remove];
384        if system_prefix >= drain_end {
385            return; // nothing to drain after system messages
386        }
387
388        self.chat_messages.drain(system_prefix..drain_end);
389    }
390
391    /// Remove the last message from `chat_messages`.
392    /// Used by the max_message_tokens safety valve to discard oversized messages.
393    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    /// Replace chat messages — only for persistence restore.
432    /// Validates message sequence before replacing.
433    ///
434    /// 仅供持久化恢复使用。调用方必须保证 messages 序列合法。
435    pub fn set_chat_messages(&mut self, messages: Vec<ChatMessage>) -> Result<(), String> {
436        validate_message_sequence(&messages)?;
437        // Recalculate total_tool_calls from the incoming messages so middleware
438        // decisions (e.g. first_turn_only enforcement) see the correct count.
439        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
454/// Validate that a chat message sequence is well-formed for LLM API consumption.
455///
456/// Checks:
457/// - No Tool message without a preceding Assistant with matching tool_call
458/// - No duplicate Tool messages for the same tool_call_id
459/// - All tool_calls in an Assistant batch must be answered before the next Assistant batch
460/// - No unanswered tool calls at the end of the sequence
461pub 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                // Remove the ID on match — also detects duplicates (second remove returns false)
474                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                // Previous batch must be fully answered before a new batch starts
486                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    // All tool calls must be answered by the end of the sequence
499    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        // One user turn: User -> Assistant(tool_calls) -> Tool -> Assistant(text)
545        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        // Turn 1
563        s.push_message(MessageRole::User, "u1");
564        s.push_message(MessageRole::Assistant, "a1");
565        // Turn 2
566        s.push_message(MessageRole::User, "u2");
567        s.push_message(MessageRole::Assistant, "a2");
568        // Turn 3
569        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        // System message preserved
575        assert!(matches!(s.chat_messages()[0], ChatMessage::System { .. }));
576        // Oldest user message is u2
577        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        // Turn 1 with tool call
586        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        // Turn 2
591        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        // chat_messages should have lost 4 entries (User, Assistant(tool), Tool, Assistant(text))
599        assert_eq!(s.chat_messages().len(), chat_before - 4);
600        // simple_messages (derived from chat_messages, tool_calls-only filtered) loses 3 entries
601        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); // only User in simple_messages (tool_calls-only filtered)
622        s.pop_last_message();
623        assert_eq!(s.chat_messages().len(), 1);
624        assert_eq!(s.simple_messages().len(), 1); // simple_messages unchanged (still just User)
625    }
626
627    #[test]
628    fn test_pop_last_message_empty_session() {
629        let mut s = make_session();
630        s.pop_last_message(); // should not panic
631        assert_eq!(s.chat_messages().len(), 0);
632    }
633
634    // ── B5: remaining session lifecycle paths ──────────────────────────────
635
636    #[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"); // only c1 answered
763        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    // ── RunState tests ────────────────────────────────────────────────────
884
885    #[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        // original_thinking_enabled is NOT reset
916        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); // reset
932        assert_eq!(rs.empty_response_strikes, 0); // reset
933    }
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); // reset
957    }
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); // reset
980    }
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        // Simulate: reasoning only → tool calls → reasoning only
994        let mut rs = RunState::default();
995
996        // Branch 1: reasoning only
997        rs.record_reasoning_only();
998        assert_eq!(rs.reasoning_only_strikes, 1);
999
1000        // Branch 3: tool calls (should reset reasoning_only_strikes)
1001        rs.record_tool_calls(2);
1002        assert_eq!(rs.reasoning_only_strikes, 0);
1003        assert_eq!(rs.turn_tool_calls, 2);
1004
1005        // Branch 1 again: reasoning only (should start from 1, not 2)
1006        let strikes = rs.record_reasoning_only();
1007        assert_eq!(strikes, 1);
1008    }
1009
1010    #[test]
1011    fn run_state_empty_to_reasoning_reset() {
1012        // Simulate: empty → empty → reasoning only (should reset empty strikes)
1013        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        // Branch 1: reasoning only (should reset empty_response_strikes)
1020        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        // After 1st reasoning-only: not disabled
1036        rs.record_reasoning_only();
1037        assert!(!rs.thinking_disabled_for_rest_of_run);
1038        assert_eq!(rs.reasoning_only_strikes, 1);
1039
1040        // After 2nd reasoning-only: not disabled
1041        rs.record_reasoning_only();
1042        assert!(!rs.thinking_disabled_for_rest_of_run);
1043        assert_eq!(rs.reasoning_only_strikes, 2);
1044
1045        // After 3rd reasoning-only: disabled!
1046        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        // Simulate 3 reasoning-only responses
1056        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        // Reset for new run (new user message)
1063        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        // Simulate 3 reasoning-only responses
1073        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        // Tool calls should NOT reset thinking_disabled_for_rest_of_run
1079        // (it should stay disabled for the rest of the run)
1080        rs.record_tool_calls(2);
1081        assert!(rs.thinking_disabled_for_rest_of_run);
1082        assert_eq!(rs.reasoning_only_strikes, 0); // strikes reset, but thinking stays disabled
1083    }
1084
1085    // ── Backward-compatible deserialization ────────────────────────────────
1086
1087    #[test]
1088    fn deserialize_legacy_flat_fields() {
1089        // Old format: nudge_count, turn_tool_calls, etc. as flat fields
1090        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); // default
1106    }
1107
1108    #[test]
1109    fn deserialize_new_run_state_format() {
1110        // New format: nested run_state
1111        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        // When both are present, run_state wins
1134        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); // run_state wins
1151        assert_eq!(session.run_state.nudge_count, 0); // run_state wins
1152    }
1153
1154    #[test]
1155    fn deserialize_legacy_missing_optional_fields() {
1156        // Old format with some fields missing (defaults to 0)
1157        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); // missing → 0
1167        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        // Valid JSON args should pass through unchanged
1192        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        // Truncated (invalid JSON) args should be wrapped in an error object
1209        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            // The arguments should now be valid JSON (the error wrapper)
1221            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        // Bug-2: Invalid JSON args with multi-byte UTF-8 chars near byte 200
1233        // cause a panic at char boundary when slicing &args[..args.len().min(200)].
1234        let mut s = make_session();
1235
1236        // Build invalid JSON with CJK chars that straddle the 200-byte boundary.
1237        // "あ" = 3 bytes in UTF-8. Repeating ~70 times = ~210 bytes, then add invalid suffix.
1238        let mut bad_args = "あ".repeat(70); // 70 * 3 = 210 bytes
1239        bad_args.push_str("truncated"); // makes it invalid JSON
1240
1241        // This must NOT panic — the preview slice should respect char boundaries.
1242        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            // Should be wrapped in error object (valid JSON)
1250            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        // Regression test for: when LLM response is truncated (finish_reason=max_tokens),
1261        // the code must push assistant message WITH tool_use blocks (not plain text)
1262        // so that subsequent tool_result messages can match the tool_use_id.
1263        // This is required by Anthropic protocol.
1264        let mut s = make_session();
1265
1266        // Simulate truncated tool call response
1267        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        // Push assistant message with tool_calls (as the fix does)
1281        s.push_assistant_tool_calls(
1282            &tool_calls,
1283            Some("thinking...".to_string()),
1284            Some("I'll help you".to_string()),
1285        );
1286
1287        // Push tool results for each tool call
1288        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        // Verify the message sequence is valid
1296        // 1. Assistant message should have tool_calls
1297        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        // 2. Tool result messages should have matching tool_use_id
1312        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        // 3. Validate message sequence (this would catch the bug)
1325        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        // ── RunState property tests ─────────────────────────────────────────
1339
1340        #[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            // original_thinking_enabled is NOT reset
1365            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        // ── push_assistant_tool_calls property tests ────────────────────────
1399
1400        #[test]
1401        fn push_assistant_tool_calls_valid_json_unchanged(args in r"\{[^{}]{0,200}\}") {
1402            // Only test strings that are actually valid JSON objects
1403            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            // Skip if it happens to be valid JSON
1424            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        // ── trim_oldest_turns property tests ────────────────────────────────
1443
1444        #[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        // ── validate_message_sequence property tests ────────────────────────
1456
1457        #[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}