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