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    /// Number of consecutive LLM turns where the truncation guard fired
60    /// (all tool_calls had invalid/truncated arguments). Reset to 0 when
61    /// at least one valid tool call executes successfully. Used to break
62    /// the re-issue death spiral: after N consecutive truncations, stop
63    /// retrying and fail with an actionable message.
64    #[serde(default)]
65    pub truncation_strikes: usize,
66}
67
68impl RunState {
69    /// Reset all run-level state for a new run (when a new user message arrives).
70    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        // Note: original_thinking_enabled is NOT reset here
79        // It should be set once when the session starts
80    }
81
82    /// Record tool calls (branch 3: tool calls).
83    /// Resets reasoning_only_strikes, empty_response_strikes, and truncation_strikes.
84    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    /// Record reasoning-only response (branch 1).
93    /// Resets empty_response_strikes.
94    /// Returns the new strike count.
95    /// If strikes reach 3, disables thinking for the rest of the run.
96    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    /// Record empty response (branch 2).
110    /// Resets reasoning_only_strikes.
111    /// Returns the new strike count.
112    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    /// Record a truncation guard firing (all tool_calls had invalid arguments).
119    /// Returns the new strike count. Caller should check against threshold.
120    pub fn record_truncation(&mut self) -> usize {
121        self.truncation_strikes += 1;
122        self.truncation_strikes
123    }
124}
125
126/// Deserialization shim for backward compatibility.
127///
128/// Before `RunState` was introduced, `nudge_count`, `turn_tool_calls`,
129/// `reasoning_only_strikes`, and `empty_response_strikes` were flat fields
130/// on `AgentSession`.  This struct accepts **both** the old flat format and
131/// the new nested `run_state` format, migrating legacy data on the fly.
132#[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    // ── new format (preferred) ──
140    run_state: Option<RunState>,
141
142    // ── legacy flat fields (fallback) ──
143    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/// Stable identity + rolling state of a single chat thread.
169///
170/// `Default` returns a fresh session; load persisted state via serde.
171/// Backward-compatible with the old flat-field format (pre-`RunState` migration).
172#[derive(Clone, Debug, Default, Serialize, Deserialize)]
173#[serde(from = "RawAgentSession")]
174pub struct AgentSession {
175    id: Option<SessionId>,
176    /// LLM API format messages, sent directly to the provider.
177    /// This is the single source of truth for the conversation state.
178    chat_messages: Vec<ChatMessage>,
179    always_allowed_actions: HashSet<String>,
180    /// Total number of tool calls made in this session (across all turns).
181    /// Used by middleware for decisions like "first_turn_only" enforcement.
182    pub total_tool_calls: usize,
183    /// Run-level state tracking for the react loop.
184    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    /// Derive a simplified `Vec<Message>` view from the canonical `chat_messages`.
203    /// Assistant messages that contain only tool_calls (no text content) are
204    /// skipped, since they have no corresponding simplified representation.
205    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    /// 可变引用,仅用于需要直接操作消息的高级场景。
225    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    /// Push an assistant message with reasoning/thinking content preserved.
249    /// This allows the LLM to see its own prior reasoning in subsequent turns,
250    /// preventing it from re-deriving the same conclusions every turn.
251    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                // Invalid JSON args (provider truncated the stream mid-generation)
292                // are sanitized to an empty object. We deliberately do NOT wrap
293                // them in a descriptive `{error, original_args_preview, message}`
294                // object: "{}" is valid JSON so the next request is never rejected
295                // with 400, and — critically — a rich error object here becomes an
296                // imitation vector: the model replays it verbatim as the next
297                // call's arguments and, because it parses, it slips past the react
298                // truncation guard. The truncation explanation lives in the
299                // paired tool_result, which is where the model reads feedback.
300                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    /// 移除所有临时消息(ephemeral=true)。
332    ///
333    /// 在 turn 结束时调用,确保注入的临时内容不残留到下一轮。
334    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    /// Count the number of conversation turns.
348    /// A turn starts with a User message and includes subsequent Assistant/Tool messages.
349    pub fn turn_count(&self) -> usize {
350        self.chat_messages
351            .iter()
352            .filter(|m| matches!(m, ChatMessage::User { .. }))
353            .count()
354    }
355
356    /// Remove the oldest turns from the front until turn count ≤ max_turns.
357    /// Preserves the System message at index 0 if present.
358    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        // Find User message positions (turn boundaries) in chat_messages
366        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        // Preserve system prefix: count leading System messages
384        let system_prefix = self
385            .chat_messages
386            .iter()
387            .take_while(|m| matches!(m, ChatMessage::System { .. }))
388            .count();
389
390        // Drain from system_prefix up to the start of the (turns_to_remove + 1)-th turn
391        let drain_end = user_positions[turns_to_remove];
392        if system_prefix >= drain_end {
393            return; // nothing to drain after system messages
394        }
395
396        self.chat_messages.drain(system_prefix..drain_end);
397    }
398
399    /// Remove the last message from `chat_messages`.
400    /// Used by the max_message_tokens safety valve to discard oversized messages.
401    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    /// Replace chat messages — only for persistence restore.
440    /// Validates message sequence before replacing.
441    ///
442    /// 仅供持久化恢复使用。调用方必须保证 messages 序列合法。
443    pub fn set_chat_messages(&mut self, messages: Vec<ChatMessage>) -> Result<(), String> {
444        validate_message_sequence(&messages)?;
445        // Recalculate total_tool_calls from the incoming messages so middleware
446        // decisions (e.g. first_turn_only enforcement) see the correct count.
447        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
462/// Validate that a chat message sequence is well-formed for LLM API consumption.
463///
464/// Checks:
465/// - No Tool message without a preceding Assistant with matching tool_call
466/// - No duplicate Tool messages for the same tool_call_id
467/// - All tool_calls in an Assistant batch must be answered before the next Assistant batch
468/// - No unanswered tool calls at the end of the sequence
469pub 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                // Remove the ID on match — also detects duplicates (second remove returns false)
482                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                // Previous batch must be fully answered before a new batch starts
494                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    // All tool calls must be answered by the end of the sequence
507    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        // One user turn: User -> Assistant(tool_calls) -> Tool -> Assistant(text)
553        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        // Turn 1
571        s.push_message(MessageRole::User, "u1");
572        s.push_message(MessageRole::Assistant, "a1");
573        // Turn 2
574        s.push_message(MessageRole::User, "u2");
575        s.push_message(MessageRole::Assistant, "a2");
576        // Turn 3
577        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        // System message preserved
583        assert!(matches!(s.chat_messages()[0], ChatMessage::System { .. }));
584        // Oldest user message is u2
585        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        // Turn 1 with tool call
594        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        // Turn 2
599        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        // chat_messages should have lost 4 entries (User, Assistant(tool), Tool, Assistant(text))
607        assert_eq!(s.chat_messages().len(), chat_before - 4);
608        // simple_messages (derived from chat_messages, tool_calls-only filtered) loses 3 entries
609        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); // only User in simple_messages (tool_calls-only filtered)
630        s.pop_last_message();
631        assert_eq!(s.chat_messages().len(), 1);
632        assert_eq!(s.simple_messages().len(), 1); // simple_messages unchanged (still just User)
633    }
634
635    #[test]
636    fn test_pop_last_message_empty_session() {
637        let mut s = make_session();
638        s.pop_last_message(); // should not panic
639        assert_eq!(s.chat_messages().len(), 0);
640    }
641
642    // ── B5: remaining session lifecycle paths ──────────────────────────────
643
644    #[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"); // only c1 answered
771        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    // ── RunState tests ────────────────────────────────────────────────────
892
893    #[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        // original_thinking_enabled is NOT reset
926        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); // reset
942        assert_eq!(rs.empty_response_strikes, 0); // reset
943    }
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); // reset
967    }
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); // reset
990    }
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        // Simulate: reasoning only → tool calls → reasoning only
1004        let mut rs = RunState::default();
1005
1006        // Branch 1: reasoning only
1007        rs.record_reasoning_only();
1008        assert_eq!(rs.reasoning_only_strikes, 1);
1009
1010        // Branch 3: tool calls (should reset reasoning_only_strikes)
1011        rs.record_tool_calls(2);
1012        assert_eq!(rs.reasoning_only_strikes, 0);
1013        assert_eq!(rs.turn_tool_calls, 2);
1014
1015        // Branch 1 again: reasoning only (should start from 1, not 2)
1016        let strikes = rs.record_reasoning_only();
1017        assert_eq!(strikes, 1);
1018    }
1019
1020    #[test]
1021    fn run_state_empty_to_reasoning_reset() {
1022        // Simulate: empty → empty → reasoning only (should reset empty strikes)
1023        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        // Branch 1: reasoning only (should reset empty_response_strikes)
1030        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        // After 1st reasoning-only: not disabled
1046        rs.record_reasoning_only();
1047        assert!(!rs.thinking_disabled_for_rest_of_run);
1048        assert_eq!(rs.reasoning_only_strikes, 1);
1049
1050        // After 2nd reasoning-only: not disabled
1051        rs.record_reasoning_only();
1052        assert!(!rs.thinking_disabled_for_rest_of_run);
1053        assert_eq!(rs.reasoning_only_strikes, 2);
1054
1055        // After 3rd reasoning-only: disabled!
1056        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        // Simulate 3 reasoning-only responses
1066        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        // Reset for new run (new user message)
1073        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        // Simulate 3 reasoning-only responses
1083        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        // Tool calls should NOT reset thinking_disabled_for_rest_of_run
1089        // (it should stay disabled for the rest of the run)
1090        rs.record_tool_calls(2);
1091        assert!(rs.thinking_disabled_for_rest_of_run);
1092        assert_eq!(rs.reasoning_only_strikes, 0); // strikes reset, but thinking stays disabled
1093    }
1094
1095    // ── Backward-compatible deserialization ────────────────────────────────
1096
1097    #[test]
1098    fn deserialize_legacy_flat_fields() {
1099        // Old format: nudge_count, turn_tool_calls, etc. as flat fields
1100        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); // default
1116    }
1117
1118    #[test]
1119    fn deserialize_new_run_state_format() {
1120        // New format: nested run_state
1121        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        // When both are present, run_state wins
1144        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); // run_state wins
1161        assert_eq!(session.run_state.nudge_count, 0); // run_state wins
1162    }
1163
1164    #[test]
1165    fn deserialize_legacy_missing_optional_fields() {
1166        // Old format with some fields missing (defaults to 0)
1167        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); // missing → 0
1177        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        // Valid JSON args should pass through unchanged
1202        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        // Truncated (invalid JSON) args must be sanitized to an empty object,
1219        // NOT wrapped in a descriptive error object. Two reasons:
1220        //   1. "{}" is valid JSON, so the next request is never rejected 400.
1221        //   2. a rich `{error, original_args_preview, message}` object in
1222        //      assistant history is an imitation vector — the model replays it
1223        //      verbatim as the next call's arguments, and because it parses as
1224        //      valid JSON it slips past the react truncation guard and resurfaces
1225        //      as a downstream ToolArgsInvalid. The truncation explanation
1226        //      belongs in the tool_result, not the assistant arguments.
1227        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        // Bug-2: Invalid JSON args with multi-byte UTF-8 chars near byte 200
1251        // cause a panic at char boundary when slicing &args[..args.len().min(200)].
1252        let mut s = make_session();
1253
1254        // Build invalid JSON with CJK chars that straddle the 200-byte boundary.
1255        // "あ" = 3 bytes in UTF-8. Repeating ~70 times = ~210 bytes, then add invalid suffix.
1256        let mut bad_args = "あ".repeat(70); // 70 * 3 = 210 bytes
1257        bad_args.push_str("truncated"); // makes it invalid JSON
1258
1259        // This must NOT panic. Invalid args are sanitized to "{}" (a fixed
1260        // literal), so no slicing of the multibyte string happens at all.
1261        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        // Regression test for: when LLM response is truncated (finish_reason=max_tokens),
1277        // the code must push assistant message WITH tool_use blocks (not plain text)
1278        // so that subsequent tool_result messages can match the tool_use_id.
1279        // This is required by Anthropic protocol.
1280        let mut s = make_session();
1281
1282        // Simulate truncated tool call response
1283        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        // Push assistant message with tool_calls (as the fix does)
1297        s.push_assistant_tool_calls(
1298            &tool_calls,
1299            Some("thinking...".to_string()),
1300            Some("I'll help you".to_string()),
1301        );
1302
1303        // Push tool results for each tool call
1304        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        // Verify the message sequence is valid
1312        // 1. Assistant message should have tool_calls
1313        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        // 2. Tool result messages should have matching tool_use_id
1328        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        // 3. Validate message sequence (this would catch the bug)
1341        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        // ── RunState property tests ─────────────────────────────────────────
1355
1356        #[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            // original_thinking_enabled is NOT reset
1383            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        // ── push_assistant_tool_calls property tests ────────────────────────
1417
1418        #[test]
1419        fn push_assistant_tool_calls_valid_json_unchanged(args in r"\{[^{}]{0,200}\}") {
1420            // Only test strings that are actually valid JSON objects
1421            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            // Skip if it happens to be valid JSON
1442            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                // Must be valid JSON (no 400 on replay) AND carry no poison
1453                // wrapper the model could echo back as arguments.
1454                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        // ── trim_oldest_turns property tests ────────────────────────────────
1463
1464        #[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        // ── validate_message_sequence property tests ────────────────────────
1476
1477        #[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}