1use crate::core::agent::error_recovery::ErrorRecoveryState;
4use crate::core::agent::task::{TaskOutcome, TaskResults};
5use crate::core::pending_actions::PendingActions;
6use crate::core::state_schema::SchemaVersion;
7use crate::exec::events::Usage;
8use crate::llm::provider::{Message, ResponsesContinuationState, responses_continuation_key};
9use crate::llm::providers::gemini::wire::{Content, FunctionResponse, Part};
10use crate::llm::request_gap::RequestGapTracker;
11use hashbrown::HashMap;
12use parking_lot::Mutex;
13use std::collections::VecDeque;
14use std::sync::Arc;
15use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
16use vtcode_commons::ErrorCategory;
17use vtcode_exec_events::ThreadEvent;
18
19#[derive(Debug, Clone)]
20pub struct ToolExecutionObservation {
21 pub tool_name: String,
22 pub attempts: u32,
23 pub duration_ms: u64,
24 pub error_category: Option<ErrorCategory>,
25}
26
27pub struct AgentSessionState {
30 pub session_id: String,
32
33 pub conversation: Vec<Content>,
35
36 pub messages: Arc<Vec<Message>>,
42
43 pub schema_version: SchemaVersion,
45
46 pub stats: SessionStats,
48 pub auto_compact_suppressed: u8,
52
53 pub constraints: SessionConstraints,
55
56 pub outcome: TaskOutcome,
58 pub stop_reason: Option<String>,
60 pub total_cost_usd: Option<f64>,
62
63 pub is_completed: bool,
65
66 pub current_stage: Option<String>,
68
69 pub created_contexts: Vec<String>,
71 pub modified_files: Vec<String>,
72 pub executed_commands: Vec<String>,
73 pub warnings: Vec<String>,
74 pub last_file_path: Option<String>,
75 pub last_dir_path: Option<String>,
76
77 pub consecutive_tool_loops: usize,
79 pub tool_loop_limit_hit: bool,
80 pub consecutive_escalations: u32,
83 pub progress_hashes: VecDeque<u64>,
86 pub stagnant_turns: usize,
88 pub last_processed_message_idx: usize,
89 pub previous_response_chains: HashMap<(String, String), ResponsesContinuationState>,
91 pub error_recovery: Arc<Mutex<ErrorRecoveryState>>,
93 pub pending_actions: PendingActions,
95
96 pub consecutive_idle_turns: usize,
98 pub max_tool_loop_streak: usize,
99 pub turn_count: usize,
100 pub turn_total_ms: u128,
101 pub turn_max_ms: u128,
102 pub turn_durations_ms: Vec<u128>,
103 pub turn_tool_observations: Vec<ToolExecutionObservation>,
105 cached_total_tokens: usize,
108
109 request_gap: RequestGapTracker,
114
115 last_reasoning_effort: Option<crate::config::types::ReasoningEffortLevel>,
118 last_model: Option<String>,
123}
124
125#[derive(Debug, Default, Clone)]
127pub struct SessionStats {
128 pub turns_executed: usize,
129 pub total_duration: Duration,
130 pub turn_durations: Vec<Duration>,
131 pub total_usage: Usage,
132 pub provider_name: String,
137 pub prompt_cache_health: super::cache_health::PromptCacheHealthMonitor,
141}
142
143impl SessionStats {
144 pub fn merge_usage(
148 &mut self,
149 usage: &crate::llm::provider::Usage,
150 ) -> Option<super::cache_health::CacheHealthAlert> {
151 let normalized = crate::llm::usage_cost::normalized_turn_usage(&self.provider_name, usage);
152 self.total_usage.add(&normalized);
153 self.prompt_cache_health.record_turn(&normalized)
154 }
155}
156
157#[derive(Debug, Clone)]
159pub struct SessionConstraints {
160 pub max_turns: usize,
161 pub max_tool_loops: usize,
162 pub max_context_tokens: usize,
163}
164
165impl AgentSessionState {
166 pub fn new(session_id: String, max_turns: usize, max_tool_loops: usize, max_context_tokens: usize) -> Self {
167 Self {
168 session_id,
169 schema_version: SchemaVersion::CURRENT,
170 conversation: Vec::new(),
171 messages: Arc::new(Vec::new()),
172 stats: SessionStats::default(),
173 auto_compact_suppressed: crate::compaction::SUPPRESS_NONE,
174 constraints: SessionConstraints { max_turns, max_tool_loops, max_context_tokens },
175 outcome: TaskOutcome::Unknown,
176 stop_reason: None,
177 total_cost_usd: None,
178 is_completed: false,
179 current_stage: None,
180 created_contexts: Vec::with_capacity(16),
181 modified_files: Vec::with_capacity(32),
182 executed_commands: Vec::with_capacity(64),
183 warnings: Vec::with_capacity(16),
184 last_file_path: None,
185 last_dir_path: None,
186 consecutive_tool_loops: 0,
187 tool_loop_limit_hit: false,
188 consecutive_escalations: 0,
189 progress_hashes: VecDeque::with_capacity(16),
190 stagnant_turns: 0,
191 last_processed_message_idx: 0,
192 previous_response_chains: HashMap::new(),
193 error_recovery: Arc::new(Mutex::new(ErrorRecoveryState::default())),
194 pending_actions: PendingActions::new(100),
195 consecutive_idle_turns: 0,
196 max_tool_loop_streak: 0,
197 turn_count: 0,
198 turn_total_ms: 0,
199 turn_max_ms: 0,
200 turn_durations_ms: Vec::with_capacity(max_turns),
201 turn_tool_observations: Vec::with_capacity(32),
202 cached_total_tokens: 0,
203 request_gap: RequestGapTracker::default(),
204 last_reasoning_effort: None,
205 last_model: None,
206 }
207 }
208
209 pub fn note_request_sent(&mut self) {
212 self.request_gap.note_request_sent();
213 }
214
215 pub fn cache_gap_exceeds(&self, threshold: Duration) -> Option<Duration> {
221 self.request_gap.cache_gap_exceeds(threshold)
222 }
223
224 pub fn note_reasoning_effort_change(&mut self, effort: Option<crate::config::types::ReasoningEffortLevel>) -> bool {
230 let changed = matches!(
231 (self.last_reasoning_effort, effort),
232 (Some(previous), Some(current)) if previous != current
233 );
234 self.last_reasoning_effort = effort;
235 changed
236 }
237
238 pub fn note_model_change(&mut self, model: &str) -> bool {
245 let changed = self.last_model.as_deref().is_some_and(|previous| previous != model);
246 self.last_model = Some(model.to_string());
247 changed
248 }
249
250 pub fn record_turn(&mut self, start: &Instant, recorded: &mut bool) {
252 if *recorded {
253 return;
254 }
255 let duration = start.elapsed();
256 let ms = duration.as_millis() as u64;
257
258 self.stats.turns_executed += 1;
259 self.stats.total_duration += duration;
260 self.stats.turn_durations.push(duration);
261
262 self.turn_count += 1;
264 self.turn_total_ms += ms as u128;
265 self.turn_max_ms = self.turn_max_ms.max(ms as u128);
266 self.turn_durations_ms.push(ms as u128);
267
268 *recorded = true;
269 }
270
271 pub fn finalize_outcome(&mut self, max_turns: usize) {
272 if self.outcome != TaskOutcome::Unknown {
273 return;
274 }
275 if self.tool_loop_limit_hit {
277 self.outcome =
278 TaskOutcome::tool_loop_limit_reached(self.constraints.max_tool_loops, self.consecutive_tool_loops);
279 } else if self.is_completed {
280 self.outcome = TaskOutcome::Success;
281 } else if self.stats.turns_executed >= max_turns {
282 self.outcome = TaskOutcome::turn_limit_reached(max_turns, self.stats.turns_executed);
283 }
284 }
285
286 pub fn register_tool_loop(&mut self) -> usize {
287 self.consecutive_tool_loops += 1;
288 self.max_tool_loop_streak = self.max_tool_loop_streak.max(self.consecutive_tool_loops);
289 self.consecutive_tool_loops
290 }
291
292 pub fn reset_tool_loop_guard(&mut self) {
293 self.consecutive_tool_loops = 0;
294 }
295
296 pub fn previous_response_id_for(&self, provider: &str, model: &str) -> Option<String> {
297 self.previous_response_chain_for(provider, model)
298 .map(|chain| chain.response_id.clone())
299 }
300
301 pub fn previous_response_chain_for(&self, provider: &str, model: &str) -> Option<&ResponsesContinuationState> {
302 responses_continuation_key(provider, model).and_then(|key| self.previous_response_chains.get(&key))
303 }
304
305 pub fn set_previous_response_chain(
306 &mut self,
307 provider: &str,
308 model: &str,
309 response_id: Option<&str>,
310 messages: Vec<Message>,
311 ) {
312 self.set_previous_response_chain_shared(provider, model, response_id, Arc::new(messages));
313 }
314
315 pub(crate) fn set_previous_response_chain_shared(
320 &mut self,
321 provider: &str,
322 model: &str,
323 response_id: Option<&str>,
324 messages: Arc<Vec<Message>>,
325 ) {
326 let Some(key) = responses_continuation_key(provider, model) else {
327 return;
328 };
329 let Some(response_id) = response_id.map(str::trim).filter(|value| !value.is_empty()) else {
330 self.previous_response_chains.remove(&key);
331 return;
332 };
333
334 self.previous_response_chains
335 .insert(key, ResponsesContinuationState { response_id: response_id.to_string(), messages });
336 }
337
338 pub fn clear_previous_response_chain_for(&mut self, provider: &str, model: &str) {
339 if let Some(key) = responses_continuation_key(provider, model) {
340 self.previous_response_chains.remove(&key);
341 }
342 }
343
344 pub fn clear_previous_response_chain(&mut self) {
345 self.previous_response_chains.clear();
346 }
347
348 pub fn mark_tool_loop_limit_hit(&mut self) {
349 if self.tool_loop_limit_hit {
350 return;
351 }
352 self.tool_loop_limit_hit = true;
353 self.outcome =
354 TaskOutcome::tool_loop_limit_reached(self.constraints.max_tool_loops, self.consecutive_tool_loops);
355 }
356
357 #[inline]
361 pub fn messages_mut(&mut self) -> &mut Vec<Message> {
362 Arc::make_mut(&mut self.messages)
363 }
364
365 pub fn add_user_message(&mut self, text: String) {
367 self.add_user_message_with_intent(text, None);
368 }
369
370 pub fn add_user_message_with_intent(&mut self, text: String, intent_id: Option<String>) {
373 let now = SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_secs()).unwrap_or(0);
374 let tokens = text.len().saturating_div(4); let metadata = crate::core::message_metadata::MessageMetadata::user_input(now, tokens);
376 let metadata = match intent_id {
377 Some(id) => metadata.with_intent_id(id),
378 None => metadata,
379 };
380 self.conversation.push(Content::user_text(text.as_str()));
381 let msg = Message::user(text).with_metadata(metadata);
382 let msg_tokens = msg.estimate_tokens();
384 self.cached_total_tokens = self.cached_total_tokens.saturating_add(msg_tokens);
385 self.messages_mut().push(msg);
386 }
387
388 const PROGRESS_STAGNATION_THRESHOLD: usize = 4;
390
391 fn assistant_response_hash(&self) -> Option<u64> {
393 use crate::llm::provider::{MessageContent, MessageRole};
394 let last_assistant = self.messages.iter().rev().find(|m| m.role == MessageRole::Assistant)?;
395 use std::hash::{Hash, Hasher};
396 let mut hasher = std::collections::hash_map::DefaultHasher::new();
397 match &last_assistant.content {
398 MessageContent::Text(t) => t.hash(&mut hasher),
399 MessageContent::Parts(parts) => {
400 for part in parts {
401 if let crate::llm::provider::ContentPart::Text { text, .. } = part {
402 text.hash(&mut hasher);
403 }
404 }
405 }
406 }
407 Some(hasher.finish())
408 }
409
410 pub fn record_progress_hash_and_check_stagnation(&mut self) -> bool {
412 let Some(hash) = self.assistant_response_hash() else {
413 self.stagnant_turns = 0;
414 return false;
415 };
416 if self.progress_hashes.back() == Some(&hash) {
417 self.stagnant_turns += 1;
418 } else {
419 self.stagnant_turns = 0;
420 }
421 self.progress_hashes.push_back(hash);
422 if self.progress_hashes.len() > 16 {
423 self.progress_hashes.pop_front();
424 }
425 self.stagnant_turns >= Self::PROGRESS_STAGNATION_THRESHOLD
426 }
427
428 pub fn attach_metadata_to_last(&mut self, source: &str, estimated_tokens: usize) {
431 if let Some(last) = self.messages_mut().last_mut() {
432 let now = SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_secs()).unwrap_or(0);
433 let metadata = match source {
434 "llm_response" => crate::core::message_metadata::MessageMetadata::llm_response(now, estimated_tokens),
435 "tool_result" => crate::core::message_metadata::MessageMetadata::tool_result(now, estimated_tokens),
436 "system" => crate::core::message_metadata::MessageMetadata::system(now, estimated_tokens),
437 "synthetic" => crate::core::message_metadata::MessageMetadata::synthetic(now, estimated_tokens),
438 _ => crate::core::message_metadata::MessageMetadata::user_input(now, estimated_tokens),
439 };
440 last.metadata = Some(metadata);
441 }
442 }
443
444 pub fn utilization(&self) -> f64 {
446 if self.constraints.max_context_tokens == 0 {
447 return 0.0;
448 }
449 self.total_tokens() as f64 / self.constraints.max_context_tokens as f64
450 }
451
452 #[inline]
456 pub fn total_tokens(&self) -> usize {
457 self.cached_total_tokens
458 }
459
460 pub fn reconcile_token_count(&mut self) {
464 self.cached_total_tokens = self.messages.iter().map(|m| m.estimate_tokens()).sum();
465 }
466
467 #[inline]
470 pub fn adjust_token_count(&mut self, delta: isize) {
471 if delta >= 0 {
472 self.cached_total_tokens = self
474 .cached_total_tokens
475 .saturating_add(usize::try_from(delta).expect("delta >= 0 checked above"));
476 } else {
477 self.cached_total_tokens = self.cached_total_tokens.saturating_sub(delta.unsigned_abs());
478 }
479 }
480
481 pub fn preflight_token_check(
489 &self,
490 system_prompt_tokens: usize,
491 tool_def_tokens: usize,
492 reserved_output_tokens: usize,
493 ) -> (bool, usize, usize) {
494 let budget = self.constraints.max_context_tokens.saturating_sub(reserved_output_tokens);
495 let estimated = self
496 .total_tokens()
497 .saturating_add(system_prompt_tokens)
498 .saturating_add(tool_def_tokens);
499 (estimated <= budget, estimated, budget)
500 }
501
502 pub fn find_safe_split_point(&self, preferred_split_at: usize) -> usize {
504 crate::core::agent::state::safe_history_split_point(&self.messages, self.conversation.len(), preferred_split_at)
505 }
506
507 pub fn normalize(&mut self) {
509 crate::core::agent::state::normalize_history(self.messages_mut());
510 self.reconcile_token_count();
511 }
512
513 pub fn clear_conversation_history(&mut self) {
525 self.messages_mut().clear();
526 self.conversation.clear();
527 self.cached_total_tokens = 0;
528 self.last_processed_message_idx = 0;
529 self.previous_response_chains.clear();
530 self.consecutive_tool_loops = 0;
531 self.tool_loop_limit_hit = false;
532 self.consecutive_escalations = 0;
533 self.progress_hashes.clear();
534 self.stagnant_turns = 0;
535 self.pending_actions = PendingActions::new(100);
536 self.consecutive_idle_turns = 0;
537 self.max_tool_loop_streak = 0;
538 self.turn_tool_observations.clear();
539 self.current_stage = None;
540 self.auto_compact_suppressed = crate::compaction::SUPPRESS_NONE;
541 self.error_recovery.lock().reset();
542 }
543
544 pub fn into_results(
545 self,
546 summary: String,
547 thread_events: Vec<ThreadEvent>,
548 total_duration_ms: u128,
549 ) -> TaskResults {
550 let average_turn_duration_ms = if self.turn_count > 0 {
551 Some(self.turn_total_ms as f64 / self.turn_count as f64)
552 } else {
553 None
554 };
555 let max_turn_duration_ms = if self.turn_count > 0 {
556 Some(self.turn_max_ms)
557 } else {
558 None
559 };
560
561 TaskResults {
562 created_contexts: self.created_contexts,
563 modified_files: self.modified_files,
564 executed_commands: self.executed_commands,
565 summary,
566 stop_reason: self.stop_reason,
567 total_cost_usd: self.total_cost_usd,
568 warnings: self.warnings,
569 thread_events,
570 outcome: self.outcome,
571 turns_executed: self.stats.turns_executed,
572 total_duration_ms,
573 average_turn_duration_ms,
574 max_turn_duration_ms,
575 turn_durations_ms: self.turn_durations_ms,
576 }
577 }
578
579 fn push_tool_event(&mut self, call_id: String, tool_name: &str, value: &serde_json::Value, is_gemini: bool) {
584 if is_gemini {
585 self.conversation.push(Content {
586 role: "function".to_string(),
587 parts: vec![Part::FunctionResponse {
588 function_response: FunctionResponse {
589 name: tool_name.to_string(),
590 response: value.clone(),
591 id: Some(call_id.clone()),
592 },
593 thought_signature: None,
594 }],
595 });
596 }
597 let serialized = serde_json::to_string(value).expect("Value serialization is infallible");
598 let msg = Message::tool_response(call_id, serialized);
599 let tokens = msg.estimate_tokens();
600 self.cached_total_tokens = self.cached_total_tokens.saturating_add(tokens);
601 self.messages_mut().push(msg);
602 }
603
604 pub fn push_tool_result(&mut self, call_id: String, tool_name: &str, result: &serde_json::Value, is_gemini: bool) {
606 self.push_tool_event(call_id, tool_name, result, is_gemini);
607 let tool_name_owned = tool_name.to_owned();
610 if !self.executed_commands.contains(&tool_name_owned) {
611 self.executed_commands.push(tool_name_owned);
612 }
613 }
614
615 pub fn push_tool_error(
617 &mut self,
618 call_id: String,
619 tool_name: &str,
620 error_payload: &serde_json::Value,
621 is_gemini: bool,
622 ) {
623 self.push_tool_event(call_id, tool_name, error_payload, is_gemini);
624 }
625
626 pub fn push_warning(&mut self, warning: impl Into<String>) {
637 let warning = warning.into();
638 if self.warnings.contains(&warning) {
639 return;
640 }
641 if self.warnings.len() < MAX_SESSION_WARNINGS {
642 self.warnings.push(warning);
643 } else if self.warnings.last().is_none_or(|w| w != WARNINGS_ELIDED_MARKER) {
644 self.warnings.push(WARNINGS_ELIDED_MARKER.to_string());
645 }
646 }
647}
648
649const MAX_SESSION_WARNINGS: usize = 200;
653const WARNINGS_ELIDED_MARKER: &str = "[Additional unique warnings elided — session warning cap reached]";
654
655#[cfg(test)]
656mod tests {
657 use super::AgentSessionState;
658 use crate::config::types::ReasoningEffortLevel;
659 use crate::llm::provider::Message;
660 use crate::llm::providers::gemini::wire::Part;
661 use std::sync::Arc;
662 use std::thread;
663 use std::time::Duration;
664
665 #[test]
669 fn cache_gap_exceeds_delegates_to_request_gap_tracker() {
670 let mut state = AgentSessionState::new("session".to_string(), 4, 4, 16_000);
671 assert_eq!(state.cache_gap_exceeds(Duration::from_millis(1)), None);
672
673 state.note_request_sent();
674 thread::sleep(Duration::from_millis(15));
675 let gap = state.cache_gap_exceeds(Duration::from_millis(5));
676 assert!(gap.is_some_and(|elapsed| elapsed >= Duration::from_millis(15)));
677 }
678
679 #[test]
680 fn note_reasoning_effort_change_is_false_on_first_request() {
681 let mut state = AgentSessionState::new("session".to_string(), 4, 4, 16_000);
682 assert!(!state.note_reasoning_effort_change(Some(ReasoningEffortLevel::Medium)));
683 }
684
685 #[test]
686 fn note_reasoning_effort_change_is_false_when_unchanged() {
687 let mut state = AgentSessionState::new("session".to_string(), 4, 4, 16_000);
688 assert!(!state.note_reasoning_effort_change(Some(ReasoningEffortLevel::Medium)));
689 assert!(!state.note_reasoning_effort_change(Some(ReasoningEffortLevel::Medium)));
690 }
691
692 #[test]
693 fn note_reasoning_effort_change_is_true_when_changed() {
694 let mut state = AgentSessionState::new("session".to_string(), 4, 4, 16_000);
695 assert!(!state.note_reasoning_effort_change(Some(ReasoningEffortLevel::Medium)));
696 assert!(state.note_reasoning_effort_change(Some(ReasoningEffortLevel::High)));
697 assert!(!state.note_reasoning_effort_change(Some(ReasoningEffortLevel::High)));
699 }
700
701 #[test]
702 fn note_model_change_is_false_on_first_request() {
703 let mut state = AgentSessionState::new("session".to_string(), 4, 4, 16_000);
704 assert!(!state.note_model_change("gpt-5.6"));
705 }
706
707 #[test]
708 fn note_model_change_is_false_when_unchanged() {
709 let mut state = AgentSessionState::new("session".to_string(), 4, 4, 16_000);
710 assert!(!state.note_model_change("gpt-5.6"));
711 assert!(!state.note_model_change("gpt-5.6"));
712 }
713
714 #[test]
715 fn note_model_change_is_true_when_changed() {
716 let mut state = AgentSessionState::new("session".to_string(), 4, 4, 16_000);
717 assert!(!state.note_model_change("gpt-5.6"));
718 assert!(state.note_model_change("claude-opus-4-6"));
719 assert!(!state.note_model_change("claude-opus-4-6"));
721 }
722
723 #[test]
724 fn previous_response_chain_is_scoped_to_provider_and_model() {
725 let mut state = AgentSessionState::new("session".to_string(), 4, 4, 16_000);
726 let messages_52 = vec![Message::user("hello".to_string())];
727 let messages_54 = vec![Message::user("continue".to_string())];
728
729 state.set_previous_response_chain("openai", "gpt-5.6", Some("resp_123"), messages_52.clone());
730 state.set_previous_response_chain("openai", "gpt-5.6-sol", Some("resp_456"), messages_54.clone());
731
732 assert_eq!(state.previous_response_id_for("openai", "gpt-5.6"), Some("resp_123".to_string()));
733 assert_eq!(state.previous_response_id_for("openai", "gpt-5.6-sol"), Some("resp_456".to_string()));
734 assert_eq!(state.previous_response_id_for("gemini", "gpt-5.6"), None);
735
736 state.clear_previous_response_chain_for("openai", "gpt-5.6");
737
738 assert_eq!(state.previous_response_id_for("openai", "gpt-5.6"), None);
739 assert_eq!(state.previous_response_chain_for("openai", "gpt-5.6"), None);
740 assert_eq!(state.previous_response_id_for("openai", "gpt-5.6-sol"), Some("resp_456".to_string()));
741 assert_eq!(
742 state
743 .previous_response_chain_for("openai", "gpt-5.6-sol")
744 .map(|chain| chain.messages.as_slice()),
745 Some(messages_54.as_slice())
746 );
747
748 state.clear_previous_response_chain();
749 assert_eq!(state.previous_response_id_for("openai", "gpt-5.6-sol"), None);
750 assert_eq!(state.previous_response_chain_for("openai", "gpt-5.6-sol"), None);
751 }
752
753 #[test]
754 fn shared_previous_response_chain_reuses_history_arc() {
755 let mut state = AgentSessionState::new("session".to_string(), 4, 4, 16_000);
756 let messages = Arc::new(vec![Message::user("hello".to_string())]);
757
758 state.set_previous_response_chain_shared("openai", "gpt-5.6-sol", Some("resp_123"), Arc::clone(&messages));
759
760 let stored = &state
761 .previous_response_chain_for("openai", "gpt-5.6-sol")
762 .expect("continuation state")
763 .messages;
764 assert!(Arc::ptr_eq(&messages, stored));
765 }
766
767 #[test]
768 fn register_tool_loop_tracks_current_and_max_streak() {
769 let mut state = AgentSessionState::new("session".to_string(), 4, 4, 16_000);
770
771 assert_eq!(state.register_tool_loop(), 1);
772 assert_eq!(state.register_tool_loop(), 2);
773 assert_eq!(state.consecutive_tool_loops, 2);
774 assert_eq!(state.max_tool_loop_streak, 2);
775
776 state.reset_tool_loop_guard();
777 assert_eq!(state.register_tool_loop(), 1);
778 assert_eq!(state.max_tool_loop_streak, 2);
779 }
780
781 #[test]
782 fn push_tool_error_preserves_structured_json_for_gemini() {
783 let mut state = AgentSessionState::new("session".to_string(), 4, 4, 16_000);
784 let payload = serde_json::json!({
785 "error": {
786 "tool_name": "read_file",
787 "message": "missing file",
788 "category": "ResourceNotFound"
789 }
790 });
791
792 state.push_tool_error("call_1".to_string(), "read_file", &payload, true);
793
794 match &state.conversation[0].parts[0] {
795 Part::FunctionResponse { function_response, .. } => {
796 assert_eq!(function_response.response["error"]["message"], "missing file");
797 }
798 other => panic!("expected function response, got {other:?}"),
799 }
800 let expected_serialized = serde_json::to_string(&payload).unwrap();
801 assert_eq!(state.messages[0], Message::tool_response("call_1".to_string(), expected_serialized));
802 }
803
804 #[test]
805 fn cached_total_tokens_matches_direct_computation() {
806 let mut state = AgentSessionState::new("session".to_string(), 4, 4, 16_000);
807
808 state.add_user_message("Hello, how are you?".to_string());
810 state.push_tool_result(
811 "call_1".to_string(),
812 "read_file",
813 &serde_json::json!({"content": "test file content"}),
814 false,
815 );
816 state.push_tool_error(
817 "call_2".to_string(),
818 "write_file",
819 &serde_json::json!({"error": "permission denied"}),
820 false,
821 );
822
823 let direct = state.messages.iter().map(|m| m.estimate_tokens()).sum::<usize>();
825 assert_eq!(state.total_tokens(), direct);
826 assert!(state.total_tokens() > 0);
827 }
828
829 #[test]
830 fn reconcile_token_count_resyncs_after_external_mutation() {
831 let mut state = AgentSessionState::new("session".to_string(), 4, 4, 16_000);
832 state.add_user_message("test message".to_string());
833 let before = state.total_tokens();
834
835 state.messages_mut().push(Message::assistant("extra response".to_string()));
837 assert_ne!(state.total_tokens(), before + Message::assistant("extra response".to_string()).estimate_tokens());
838
839 state.reconcile_token_count();
841 let expected = state.messages.iter().map(|m| m.estimate_tokens()).sum::<usize>();
842 assert_eq!(state.total_tokens(), expected);
843 }
844
845 #[test]
846 fn clear_conversation_history_resets_all_state() {
847 let mut state = AgentSessionState::new("session".to_string(), 4, 4, 16_000);
848 state.messages_mut().push(Message::user("hello".to_string()));
849 state.messages_mut().push(Message::assistant("hi".to_string()));
850 state.conversation.push(crate::llm::providers::gemini::wire::Content {
851 role: "user".to_string(),
852 parts: vec![Part::Text { text: "hello".to_string(), thought_signature: None }],
853 });
854 state.reconcile_token_count();
855 state.last_processed_message_idx = 2;
856 state.progress_hashes.push_back(123);
857 state.stagnant_turns = 3;
858 state.set_previous_response_chain("openai", "gpt-5", Some("resp_1"), vec![]);
859
860 assert!(!state.messages.is_empty());
861 assert!(!state.conversation.is_empty());
862 assert!(state.total_tokens() > 0);
863
864 state.clear_conversation_history();
865
866 assert!(state.messages.is_empty());
867 assert!(state.conversation.is_empty());
868 assert_eq!(state.total_tokens(), 0);
869 assert_eq!(state.last_processed_message_idx, 0);
870 assert!(state.progress_hashes.is_empty());
871 assert_eq!(state.stagnant_turns, 0);
872 assert!(state.previous_response_chains.is_empty());
873 }
874
875 #[test]
876 fn push_warning_skips_exact_duplicates() {
877 let mut state = AgentSessionState::new("session".to_string(), 4, 4, 16_000);
878 state.push_warning("rate limited");
879 state.push_warning("rate limited");
880 state.push_warning("rate limited");
881 assert_eq!(state.warnings.len(), 1);
882 assert_eq!(state.warnings[0], "rate limited");
883 }
884
885 #[test]
886 fn push_warning_retains_distinct_warnings() {
887 let mut state = AgentSessionState::new("session".to_string(), 4, 4, 16_000);
888 state.push_warning("first warning");
889 state.push_warning("second warning");
890 state.push_warning("third warning");
891 assert_eq!(state.warnings.len(), 3);
892 assert_eq!(state.warnings[0], "first warning");
893 assert_eq!(state.warnings[1], "second warning");
894 assert_eq!(state.warnings[2], "third warning");
895 }
896
897 #[test]
898 fn push_warning_inserts_elision_marker_at_cap() {
899 let mut state = AgentSessionState::new("session".to_string(), 4, 4, 16_000);
900 for i in 0..super::MAX_SESSION_WARNINGS {
902 state.push_warning(format!("warning {i}"));
903 }
904 assert_eq!(state.warnings.len(), super::MAX_SESSION_WARNINGS);
905 assert!(!state.warnings.iter().any(|w| w == super::WARNINGS_ELIDED_MARKER));
906
907 state.push_warning("overflow A");
909 assert_eq!(state.warnings.len(), super::MAX_SESSION_WARNINGS + 1);
910 assert_eq!(state.warnings.last().unwrap(), super::WARNINGS_ELIDED_MARKER);
911
912 state.push_warning("overflow B");
915 state.push_warning("overflow C");
916 assert_eq!(state.warnings.len(), super::MAX_SESSION_WARNINGS + 1);
917 assert_eq!(state.warnings.last().unwrap(), super::WARNINGS_ELIDED_MARKER);
918 assert!(!state.warnings.iter().any(|w| w == "overflow A"));
920 assert!(!state.warnings.iter().any(|w| w == "overflow B"));
921 }
922}