Skip to main content

distri_types/
execution.rs

1use schemars::JsonSchema;
2use serde::{Deserialize, Serialize};
3use serde_json::json;
4
5use crate::{Part, PlanStep, TaskStatus, ToolResponse, core::FileType};
6
7/// Execution strategy types
8#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
9pub enum ExecutionType {
10    Interleaved,
11    Retriable,
12    React,
13    Code,
14}
15
16/// Execution result with detailed information
17#[derive(Debug, Clone, JsonSchema, Serialize, Deserialize)]
18#[serde(rename_all = "snake_case")]
19pub struct ExecutionResult {
20    pub step_id: String,
21    pub parts: Vec<Part>,
22    pub status: ExecutionStatus,
23    pub reason: Option<String>, // for rejection or failure
24    pub timestamp: i64,
25}
26
27impl ExecutionResult {
28    pub fn is_success(&self) -> bool {
29        self.status == ExecutionStatus::Success || self.status == ExecutionStatus::InputRequired
30    }
31    pub fn is_failed(&self) -> bool {
32        self.status == ExecutionStatus::Failed
33    }
34    pub fn is_rejected(&self) -> bool {
35        self.status == ExecutionStatus::Rejected
36    }
37    pub fn is_input_required(&self) -> bool {
38        self.status == ExecutionStatus::InputRequired
39    }
40
41    pub fn as_observation(&self) -> String {
42        const MAX_DATA_CHARS: usize = 500;
43        const MAX_TEXT_CHARS: usize = 1000;
44
45        // Phase 6.4: Empty result guard — prevents model issues with empty tool results
46        let has_content = self.parts.iter().any(|p| match p {
47            Part::Text(t) => !t.trim().is_empty(),
48            _ => true,
49        });
50        if !has_content && self.reason.is_none() {
51            return format!("({} completed with no output)", self.step_id);
52        }
53
54        let mut txt = String::new();
55        if let Some(reason) = &self.reason {
56            txt.push_str(reason);
57        }
58        let parts_txt = self
59            .parts
60            .iter()
61            .map(|p| match p {
62                Part::Text(text) => {
63                    if text.len() > MAX_TEXT_CHARS {
64                        let truncated: String = text.chars().take(MAX_TEXT_CHARS).collect();
65                        format!("{}... [truncated, {} total chars]", truncated, text.len())
66                    } else {
67                        text.clone()
68                    }
69                }
70                Part::ToolCall(tool_call) => format!(
71                    "Action: {} with {}",
72                    tool_call.tool_name,
73                    serde_json::to_string(&tool_call.input).unwrap_or_default()
74                ),
75                Part::Data(data) => {
76                    let serialized = serde_json::to_string(&data).unwrap_or_default();
77                    if serialized.len() > MAX_DATA_CHARS {
78                        let truncated: String = serialized.chars().take(MAX_DATA_CHARS).collect();
79                        format!(
80                            "{}... [truncated, {} total chars]",
81                            truncated,
82                            serialized.len()
83                        )
84                    } else {
85                        serialized
86                    }
87                }
88                Part::ToolResult(tool_result) => {
89                    let serialized =
90                        serde_json::to_string(&tool_result.result()).unwrap_or_default();
91                    if serialized.len() > MAX_DATA_CHARS {
92                        let truncated: String = serialized.chars().take(MAX_DATA_CHARS).collect();
93                        format!(
94                            "{}... [truncated, {} total chars]",
95                            truncated,
96                            serialized.len()
97                        )
98                    } else {
99                        serialized
100                    }
101                }
102                Part::Image(image) => match image {
103                    FileType::Url { url, .. } => format!("[Image: {}]", url),
104                    FileType::Bytes {
105                        name, mime_type, ..
106                    } => format!(
107                        "[Image: {} ({})]",
108                        name.as_deref().unwrap_or("unnamed"),
109                        mime_type
110                    ),
111                },
112                // Phase 6.2: Include artifact preview in observation
113                Part::Artifact(artifact) => {
114                    let preview = artifact
115                        .preview
116                        .as_deref()
117                        .map(|p| format!("\nPreview:\n{}", p))
118                        .unwrap_or_default();
119                    let stats_info = artifact
120                        .stats
121                        .as_ref()
122                        .map(|s| format!("{} — ", s.context_info()))
123                        .unwrap_or_default();
124                    format!(
125                        "[Artifact: {}{}\n... ({}use artifact tools for full content)]",
126                        artifact.file_id, preview, stats_info
127                    )
128                }
129            })
130            .collect::<Vec<_>>()
131            .join("\n");
132        if !parts_txt.is_empty() {
133            txt.push('\n');
134            txt.push_str(&parts_txt);
135        }
136        txt
137    }
138
139    /// Compact execution results before storing in scratchpad/history used for prompt construction.
140    ///
141    /// This keeps high-signal fields (tool ids/status/artifact refs) while stripping or truncating
142    /// large payloads that would otherwise bloat subsequent model calls.
143    pub fn compact_for_history(&self) -> Self {
144        const MAX_TEXT_CHARS: usize = 2_000;
145        const MAX_JSON_CHARS: usize = 4_000;
146
147        fn truncate(value: &str, max: usize) -> String {
148            if value.chars().count() <= max {
149                return value.to_string();
150            }
151
152            let truncated: String = value.chars().take(max).collect();
153            format!(
154                "{}\n...[truncated {} chars for history]",
155                truncated,
156                value.chars().count().saturating_sub(max)
157            )
158        }
159
160        fn compact_json(value: &serde_json::Value, max: usize) -> serde_json::Value {
161            match serde_json::to_string(value) {
162                Ok(serialized) if serialized.chars().count() > max => json!({
163                    "summary": "JSON payload omitted from history due to size",
164                    "preview": truncate(&serialized, std::cmp::min(500, max)),
165                    "truncated": true,
166                    "original_chars": serialized.chars().count()
167                }),
168                Ok(_) => value.clone(),
169                Err(_) => {
170                    json!({ "summary": "JSON payload omitted from history (serialization failed)" })
171                }
172            }
173        }
174
175        let compacted_parts = self
176            .parts
177            .iter()
178            .map(|part| match part {
179                Part::Text(text) => Part::Text(truncate(text, MAX_TEXT_CHARS)),
180                Part::Data(data) => Part::Data(compact_json(data, MAX_JSON_CHARS)),
181                Part::ToolCall(tool_call) => {
182                    let mut compacted_call = tool_call.clone();
183                    compacted_call.input = compact_json(&tool_call.input, MAX_JSON_CHARS);
184                    Part::ToolCall(compacted_call)
185                }
186                Part::ToolResult(tool_result) => {
187                    let filtered = tool_result.filter_for_save();
188                    let compacted_tool_parts = filtered
189                        .parts
190                        .iter()
191                        .map(|tool_part| match tool_part {
192                            Part::Text(text) => Part::Text(truncate(text, MAX_TEXT_CHARS)),
193                            Part::Data(data) => Part::Data(compact_json(data, MAX_JSON_CHARS)),
194                            // Keep artifact references; drop inline images from rolling context.
195                            Part::Image(_) => Part::Text(
196                                "[Image omitted from history; use artifact/reference if needed]"
197                                    .to_string(),
198                            ),
199                            other => other.clone(),
200                        })
201                        .collect();
202
203                    Part::ToolResult(ToolResponse {
204                        tool_call_id: filtered.tool_call_id,
205                        tool_name: filtered.tool_name,
206                        parts: compacted_tool_parts,
207                        parts_metadata: None,
208                    })
209                }
210                Part::Image(_) => {
211                    Part::Text("[Image omitted from history to reduce context size]".to_string())
212                }
213                Part::Artifact(artifact) => Part::Artifact(artifact.clone()),
214            })
215            .collect();
216
217        Self {
218            step_id: self.step_id.clone(),
219            parts: compacted_parts,
220            status: self.status.clone(),
221            reason: self.reason.as_ref().map(|r| truncate(r, MAX_TEXT_CHARS)),
222            timestamp: self.timestamp,
223        }
224    }
225
226    /// Maximum tokens for a single tool result in the scratchpad.
227    pub const MAX_TOOL_RESULT_TOKENS: usize = 500;
228
229    /// Ensure the result has at least one part. If empty, injects a "[No output]" guard.
230    pub fn with_empty_guard(mut self) -> Self {
231        if self.parts.is_empty() {
232            self.parts.push(Part::Text("[No output]".to_string()));
233        }
234        self
235    }
236
237    /// Compact for storage: applies `compact_for_history()` + `with_empty_guard()`.
238    pub fn compact_for_storage(&self) -> Self {
239        self.compact_for_history().with_empty_guard()
240    }
241}
242
243#[derive(Debug, Clone, JsonSchema, Serialize, Deserialize, PartialEq, Eq)]
244#[serde(rename_all = "snake_case")]
245pub enum ExecutionStatus {
246    Success,
247    Failed,
248    Rejected,
249    InputRequired,
250}
251
252impl From<ExecutionStatus> for TaskStatus {
253    fn from(val: ExecutionStatus) -> Self {
254        match val {
255            ExecutionStatus::Success => TaskStatus::Completed,
256            ExecutionStatus::Failed => TaskStatus::Failed,
257            ExecutionStatus::Rejected => TaskStatus::Canceled,
258            ExecutionStatus::InputRequired => TaskStatus::InputRequired,
259        }
260    }
261}
262
263pub enum ToolResultWithSkip {
264    ToolResult(ToolResponse),
265    // Skip tool call if it is external
266    Skip {
267        tool_call_id: String,
268        reason: String,
269    },
270}
271
272pub fn from_tool_results(tool_results: Vec<ToolResultWithSkip>) -> Vec<Part> {
273    tool_results
274        .iter()
275        .filter_map(|result| match result {
276            ToolResultWithSkip::ToolResult(tool_result) => {
277                // Simply extract parts from the tool response
278                Some(tool_result.parts.clone())
279            }
280            _ => None,
281        })
282        .flatten()
283        .collect()
284}
285
286#[derive(Debug, Clone, Serialize, Deserialize, Default)]
287pub struct ContextUsage {
288    pub tokens: u32,
289    pub input_tokens: u32,
290    pub output_tokens: u32,
291    /// Tokens read from provider cache (e.g., Anthropic prompt caching)
292    #[serde(default)]
293    pub cached_tokens: u32,
294    pub current_iteration: usize,
295    pub context_size: ContextSize,
296    /// Model used for LLM calls in this context
297    #[serde(default)]
298    pub model: Option<String>,
299    /// Per-component token budget tracking for context optimization
300    #[serde(default)]
301    pub context_budget: ContextBudget,
302    /// Snapshot taken at the start of each step — used to compute per-step deltas
303    #[serde(default)]
304    pub step_input_start: u32,
305    #[serde(default)]
306    pub step_output_start: u32,
307    #[serde(default)]
308    pub step_cached_start: u32,
309}
310
311/// Tracks token usage by component for context optimization.
312///
313/// Each field represents the estimated token count for a specific component
314/// of the prompt. This enables:
315/// - Monitoring which components consume the most context
316/// - Triggering compaction when utilization exceeds thresholds
317/// - Informing deferred loading decisions (tools, skills)
318/// - API-side prompt caching optimization
319#[derive(Debug, Clone, Serialize, Deserialize, Default)]
320pub struct ContextBudget {
321    /// Static system prompt tokens (cacheable across sessions)
322    pub system_prompt_static_tokens: usize,
323    /// Dynamic system prompt tokens (per-session: env, memory, hooks)
324    pub system_prompt_dynamic_tokens: usize,
325    /// Tool schema tokens (full schemas for core tools)
326    pub tool_schema_tokens: usize,
327    /// Deferred tool listing tokens (name + description only)
328    pub deferred_tool_tokens: usize,
329    /// Skill listing tokens in system prompt
330    pub skill_listing_tokens: usize,
331    /// Conversation history tokens (all messages)
332    pub conversation_tokens: usize,
333    /// Tool result tokens in current turn
334    pub tool_result_tokens: usize,
335    /// Total estimated context window size for the model
336    pub context_window_size: usize,
337    /// Whether the static prompt prefix hash has changed (cache bust)
338    pub static_prefix_cache_hit: bool,
339    /// Hash of the static system prompt prefix for cache tracking
340    #[serde(default)]
341    pub static_prefix_hash: Option<String>,
342}
343
344impl ContextBudget {
345    /// Total tokens currently consumed across all components
346    pub fn total_tokens(&self) -> usize {
347        self.system_prompt_static_tokens
348            + self.system_prompt_dynamic_tokens
349            + self.tool_schema_tokens
350            + self.deferred_tool_tokens
351            + self.skill_listing_tokens
352            + self.conversation_tokens
353            + self.tool_result_tokens
354    }
355
356    /// Context utilization as a percentage (0.0 - 1.0)
357    pub fn utilization(&self) -> f64 {
358        if self.context_window_size == 0 {
359            return 0.0;
360        }
361        self.total_tokens() as f64 / self.context_window_size as f64
362    }
363
364    /// Remaining tokens available in the context window
365    pub fn remaining_tokens(&self) -> usize {
366        self.context_window_size.saturating_sub(self.total_tokens())
367    }
368
369    /// Whether context utilization exceeds the warning threshold (80%)
370    pub fn is_warning(&self) -> bool {
371        self.utilization() > 0.80
372    }
373
374    /// Whether context utilization exceeds the critical threshold (90%)
375    pub fn is_critical(&self) -> bool {
376        self.utilization() > 0.90
377    }
378
379    /// Tokens saved by deferring tools (vs loading all schemas)
380    pub fn deferred_savings(&self) -> usize {
381        // This would be set externally by comparing full vs deferred tool tokens
382        0 // Placeholder - actual savings tracked by tool resolution
383    }
384}
385
386#[derive(Debug, Clone, Serialize, Deserialize, Default)]
387pub struct ContextSize {
388    pub message_count: usize,
389    pub message_chars: usize,
390    pub message_estimated_tokens: usize,
391    pub execution_history_count: usize,
392    pub execution_history_chars: usize,
393    pub execution_history_estimated_tokens: usize,
394    pub scratchpad_chars: usize,
395    pub scratchpad_estimated_tokens: usize,
396    pub total_chars: usize,
397    pub total_estimated_tokens: usize,
398    /// Per-agent context size breakdown
399    pub agent_breakdown: std::collections::HashMap<String, AgentContextSize>,
400}
401
402#[derive(Debug, Clone, Serialize, Deserialize, Default)]
403pub struct AgentContextSize {
404    pub agent_id: String,
405    pub task_count: usize,
406    pub execution_history_count: usize,
407    pub execution_history_chars: usize,
408    pub execution_history_estimated_tokens: usize,
409    pub scratchpad_chars: usize,
410    pub scratchpad_estimated_tokens: usize,
411}
412
413/// Enriched execution history entry that includes context metadata
414#[derive(Debug, Clone, Serialize, Deserialize)]
415pub struct ExecutionHistoryEntry {
416    pub thread_id: String, // Conversation context
417    pub task_id: String,   // Individual user task/request
418    pub run_id: String,    // Specific execution strand
419    pub execution_result: ExecutionResult,
420    pub stored_at: i64, // When this was stored
421}
422
423/// Entry for scratchpad formatting
424#[derive(Debug, Clone, Serialize, Deserialize)]
425pub struct ScratchpadEntry {
426    pub timestamp: i64,
427    #[serde(flatten)]
428    pub entry_type: ScratchpadEntryType,
429    pub task_id: String,
430    #[serde(default)]
431    pub parent_task_id: Option<String>,
432    pub entry_kind: Option<String>,
433}
434
435/// Type of scratchpad entry - only for Thought/Action/Observation tracking
436#[derive(Debug, Clone, Serialize, Deserialize)]
437#[serde(rename_all = "snake_case", tag = "type", content = "data")]
438pub enum ScratchpadEntryType {
439    #[serde(rename = "task")]
440    Task(Vec<Part>),
441    #[serde(rename = "plan")]
442    PlanStep(PlanStep),
443    #[serde(rename = "execution")]
444    Execution(ExecutionHistoryEntry),
445    /// Compressed summary produced by Tier 2 (semantic) compaction
446    #[serde(rename = "summary")]
447    Summary(CompactionSummary),
448    /// Skill content re-injected after compaction
449    #[serde(rename = "skill_context")]
450    SkillContext(SkillContextEntry),
451}
452
453/// Skill content re-injected after compaction to preserve agent instructions
454#[derive(Debug, Clone, Serialize, Deserialize)]
455pub struct SkillContextEntry {
456    /// Skill identifier
457    pub skill_id: String,
458    /// Full skill content (markdown)
459    pub content: String,
460    /// Timestamp when this was re-injected
461    pub reinjected_at: i64,
462}
463
464/// Summary produced by semantic compaction of older scratchpad entries
465#[derive(Debug, Clone, Serialize, Deserialize)]
466pub struct CompactionSummary {
467    /// LLM-generated summary of compacted history
468    pub summary_text: String,
469    /// Number of entries that were summarized
470    pub entries_summarized: usize,
471    /// Timestamp range of summarized entries
472    pub from_timestamp: i64,
473    pub to_timestamp: i64,
474    /// Token count saved by this compaction
475    pub tokens_saved: usize,
476}
477
478#[cfg(test)]
479mod tests {
480    use super::*;
481    use serde_json::json;
482
483    #[test]
484    fn test_scratchpad_large_observation_issue() {
485        println!("=== TESTING LARGE DATA OBSERVATION IN SCRATCHPAD ===");
486
487        // Create a very large tool response observation (similar to search results)
488        let large_data = json!({
489            "results": (0..100).map(|i| json!({
490                "id": i,
491                "name": format!("Minister {}", i),
492                "email": format!("minister{}@gov.sg", i),
493                "portfolio": format!("Ministry of Complex Affairs {}", i),
494                "biography": format!("Very long biography text that goes on and on for minister {} with lots of details about their career, education, achievements, and political history. This is intentionally verbose to demonstrate the issue with large content in scratchpad observations.", i),
495            })).collect::<Vec<_>>()
496        });
497
498        println!(
499            "Large data size: {} bytes",
500            serde_json::to_string(&large_data).unwrap().len()
501        );
502
503        // Test 1: Direct Part::Data (BROKEN - causes scratchpad bloat)
504        let execution_result_data = ExecutionResult {
505            step_id: "test-step-1".to_string(),
506            parts: vec![Part::Data(large_data.clone())],
507            status: ExecutionStatus::Success,
508            reason: None,
509            timestamp: 1234567890,
510        };
511
512        let observation_data = execution_result_data.as_observation();
513        println!(
514            "🚨 BROKEN: Direct Part::Data observation size: {} chars",
515            observation_data.len()
516        );
517        println!(
518            "Preview (first 200 chars): {}",
519            &observation_data.chars().take(200).collect::<String>()
520        );
521
522        // Test 2: File metadata (GOOD - concise)
523        let file_metadata = crate::filesystem::FileMetadata {
524            file_id: "large-search-results.json".to_string(),
525            relative_path: "thread123/task456/large-search-results.json".to_string(),
526            size: serde_json::to_string(&large_data).unwrap().len() as u64,
527            content_type: Some("application/json".to_string()),
528            original_filename: Some("search_results.json".to_string()),
529            created_at: chrono::Utc::now(),
530            updated_at: chrono::Utc::now(),
531            checksum: Some("abc123".to_string()),
532            stats: None,
533            preview: Some("JSON search results with 100 minister entries".to_string()),
534        };
535
536        let execution_result_file = ExecutionResult {
537            step_id: "test-step-2".to_string(),
538            parts: vec![Part::Artifact(file_metadata)],
539            status: ExecutionStatus::Success,
540            reason: None,
541            timestamp: 1234567890,
542        };
543
544        let observation_file = execution_result_file.as_observation();
545        println!(
546            "āœ… GOOD: File metadata observation size: {} chars",
547            observation_file.len()
548        );
549        println!("Content: {}", observation_file);
550
551        // Demonstrate the problem
552        println!("\n=== SCRATCHPAD IMPACT ===");
553        println!(
554            "āŒ Direct approach adds {} chars to scratchpad (CAUSES LOOPS!)",
555            observation_data.len()
556        );
557        println!(
558            "āœ… File metadata adds only {} chars to scratchpad",
559            observation_file.len()
560        );
561        println!(
562            "šŸ’” Size reduction: {:.1}%",
563            (1.0 - (observation_file.len() as f64 / observation_data.len() as f64)) * 100.0
564        );
565
566        // This test shows the fix is working - observations are now truncated
567        assert!(observation_data.len() < 1000, "Large data is now truncated"); // Fixed expectation
568        assert!(
569            observation_file.len() < 300,
570            "File metadata stays reasonably concise"
571        ); // Updated for detailed format
572
573        println!("\n🚨 CONCLUSION: as_observation() needs to truncate large Part::Data!");
574    }
575
576    #[test]
577    fn test_observation_truncation_fix() {
578        println!("=== TESTING OBSERVATION TRUNCATION FIX ===");
579
580        // Test large data truncation
581        let large_data = json!({
582            "big_array": (0..200).map(|i| format!("item_{}", i)).collect::<Vec<_>>()
583        });
584
585        let execution_result = ExecutionResult {
586            step_id: "test-truncation".to_string(),
587            parts: vec![Part::Data(large_data)],
588            status: ExecutionStatus::Success,
589            reason: None,
590            timestamp: 1234567890,
591        };
592
593        let observation = execution_result.as_observation();
594        println!("Truncated observation size: {} chars", observation.len());
595        println!("Content: {}", observation);
596
597        // Should be truncated and include total char count
598        assert!(
599            observation.len() < 600,
600            "Observation should be truncated to <600 chars"
601        );
602        assert!(
603            observation.contains("truncated"),
604            "Should indicate truncation"
605        );
606        assert!(
607            observation.contains("total chars"),
608            "Should show total char count"
609        );
610
611        // Test long text truncation
612        let long_text = "This is a very long text. ".repeat(100);
613        let text_result = ExecutionResult {
614            step_id: "test-text-truncation".to_string(),
615            parts: vec![Part::Text(long_text.clone())],
616            status: ExecutionStatus::Success,
617            reason: None,
618            timestamp: 1234567890,
619        };
620
621        let text_observation = text_result.as_observation();
622        println!("Text observation size: {} chars", text_observation.len());
623        assert!(
624            text_observation.len() < 1100,
625            "Text should be truncated to ~1000 chars"
626        );
627        if long_text.len() > 1000 {
628            assert!(
629                text_observation.contains("truncated"),
630                "Long text should be truncated"
631            );
632        }
633
634        println!("āœ… Observation truncation is working!");
635    }
636
637    #[test]
638    fn test_compact_for_history_filters_save_false_and_truncates_large_parts() {
639        let mut parts_metadata = std::collections::HashMap::new();
640        parts_metadata.insert(1, crate::PartMetadata { save: false });
641
642        let tool_response = ToolResponse {
643            tool_call_id: "call-1".to_string(),
644            tool_name: "search".to_string(),
645            parts: vec![
646                Part::Data(json!({"small": "kept"})),
647                Part::Data(json!({"secret": "do not persist"})),
648            ],
649            parts_metadata: Some(parts_metadata),
650        };
651
652        let huge = "x".repeat(6_000);
653        let execution_result = ExecutionResult {
654            step_id: "step-1".to_string(),
655            parts: vec![
656                Part::Text("y".repeat(2_500)),
657                Part::Data(json!({"huge": huge})),
658                Part::ToolResult(tool_response),
659            ],
660            status: ExecutionStatus::Success,
661            reason: Some("z".repeat(2_500)),
662            timestamp: 0,
663        };
664
665        let compacted = execution_result.compact_for_history();
666
667        assert_eq!(compacted.parts.len(), 3);
668        let text = match &compacted.parts[0] {
669            Part::Text(value) => value,
670            other => panic!("unexpected part: {:?}", other),
671        };
672        assert!(text.contains("[truncated"));
673
674        let data = match &compacted.parts[1] {
675            Part::Data(value) => value,
676            other => panic!("unexpected part: {:?}", other),
677        };
678        assert_eq!(data["truncated"], json!(true));
679
680        let tool = match &compacted.parts[2] {
681            Part::ToolResult(value) => value,
682            other => panic!("unexpected part: {:?}", other),
683        };
684        // save:false part should be removed.
685        assert_eq!(tool.parts.len(), 1);
686        assert!(tool.parts_metadata.is_none());
687    }
688
689    #[test]
690    fn test_context_budget_total_tokens() {
691        let budget = ContextBudget {
692            system_prompt_static_tokens: 3000,
693            system_prompt_dynamic_tokens: 2000,
694            tool_schema_tokens: 5000,
695            deferred_tool_tokens: 200,
696            skill_listing_tokens: 500,
697            conversation_tokens: 10000,
698            tool_result_tokens: 1000,
699            context_window_size: 200_000,
700            static_prefix_cache_hit: false,
701            static_prefix_hash: None,
702        };
703
704        assert_eq!(budget.total_tokens(), 21700);
705        assert!((budget.utilization() - 0.1085).abs() < 0.001);
706        assert_eq!(budget.remaining_tokens(), 178300);
707        assert!(!budget.is_warning());
708        assert!(!budget.is_critical());
709    }
710
711    #[test]
712    fn test_context_budget_warning_threshold() {
713        let budget = ContextBudget {
714            conversation_tokens: 85000,
715            context_window_size: 100_000,
716            ..Default::default()
717        };
718        assert!(budget.is_warning());
719        assert!(!budget.is_critical());
720    }
721
722    #[test]
723    fn test_context_budget_critical_threshold() {
724        let budget = ContextBudget {
725            conversation_tokens: 95000,
726            context_window_size: 100_000,
727            ..Default::default()
728        };
729        assert!(budget.is_warning());
730        assert!(budget.is_critical());
731    }
732
733    #[test]
734    fn test_context_budget_zero_window() {
735        let budget = ContextBudget::default();
736        assert_eq!(budget.utilization(), 0.0);
737        assert_eq!(budget.remaining_tokens(), 0);
738    }
739}