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        let mut txt = String::new();
46        if let Some(reason) = &self.reason {
47            txt.push_str(reason);
48        }
49        let parts_txt = self
50            .parts
51            .iter()
52            .map(|p| match p {
53                Part::Text(text) => {
54                    if text.len() > MAX_TEXT_CHARS {
55                        let truncated: String = text.chars().take(MAX_TEXT_CHARS).collect();
56                        format!(
57                            "{}... [truncated, {} total chars]",
58                            truncated,
59                            text.len()
60                        )
61                    } else {
62                        text.clone()
63                    }
64                }
65                Part::ToolCall(tool_call) => format!(
66                    "Action: {} with {}",
67                    tool_call.tool_name,
68                    serde_json::to_string(&tool_call.input).unwrap_or_default()
69                ),
70                Part::Data(data) => {
71                    let serialized = serde_json::to_string(&data).unwrap_or_default();
72                    if serialized.len() > MAX_DATA_CHARS {
73                        let truncated: String = serialized.chars().take(MAX_DATA_CHARS).collect();
74                        format!(
75                            "{}... [truncated, {} total chars]",
76                            truncated,
77                            serialized.len()
78                        )
79                    } else {
80                        serialized
81                    }
82                }
83                Part::ToolResult(tool_result) => {
84                    let serialized =
85                        serde_json::to_string(&tool_result.result()).unwrap_or_default();
86                    if serialized.len() > MAX_DATA_CHARS {
87                        let truncated: String = serialized.chars().take(MAX_DATA_CHARS).collect();
88                        format!(
89                            "{}... [truncated, {} total chars]",
90                            truncated,
91                            serialized.len()
92                        )
93                    } else {
94                        serialized
95                    }
96                }
97                Part::Image(image) => match image {
98                    FileType::Url { url, .. } => format!("[Image: {}]", url),
99                    FileType::Bytes {
100                        name, mime_type, ..
101                    } => format!(
102                        "[Image: {} ({})]",
103                        name.as_deref().unwrap_or("unnamed"),
104                        mime_type
105                    ),
106                },
107                Part::Artifact(artifact) => format!(
108                    "[Artifact ID:{}\n You can use artifact tools to read the full content\n{}]",
109                    artifact.file_id,
110                    if let Some(stats) = &artifact.stats {
111                        format!(" ({})", stats.context_info())
112                    } else {
113                        String::new()
114                    }
115                ),
116            })
117            .collect::<Vec<_>>()
118            .join("\n");
119        if !parts_txt.is_empty() {
120            txt.push('\n');
121            txt.push_str(&parts_txt);
122        }
123        txt
124    }
125
126    /// Compact execution results before storing in scratchpad/history used for prompt construction.
127    ///
128    /// This keeps high-signal fields (tool ids/status/artifact refs) while stripping or truncating
129    /// large payloads that would otherwise bloat subsequent model calls.
130    pub fn compact_for_history(&self) -> Self {
131        const MAX_TEXT_CHARS: usize = 2_000;
132        const MAX_JSON_CHARS: usize = 4_000;
133
134        fn truncate(value: &str, max: usize) -> String {
135            if value.chars().count() <= max {
136                return value.to_string();
137            }
138
139            let truncated: String = value.chars().take(max).collect();
140            format!(
141                "{}\n...[truncated {} chars for history]",
142                truncated,
143                value.chars().count().saturating_sub(max)
144            )
145        }
146
147        fn compact_json(value: &serde_json::Value, max: usize) -> serde_json::Value {
148            match serde_json::to_string(value) {
149                Ok(serialized) if serialized.chars().count() > max => json!({
150                    "summary": "JSON payload omitted from history due to size",
151                    "preview": truncate(&serialized, std::cmp::min(500, max)),
152                    "truncated": true,
153                    "original_chars": serialized.chars().count()
154                }),
155                Ok(_) => value.clone(),
156                Err(_) => {
157                    json!({ "summary": "JSON payload omitted from history (serialization failed)" })
158                }
159            }
160        }
161
162        let compacted_parts = self
163            .parts
164            .iter()
165            .map(|part| match part {
166                Part::Text(text) => Part::Text(truncate(text, MAX_TEXT_CHARS)),
167                Part::Data(data) => Part::Data(compact_json(data, MAX_JSON_CHARS)),
168                Part::ToolCall(tool_call) => {
169                    let mut compacted_call = tool_call.clone();
170                    compacted_call.input = compact_json(&tool_call.input, MAX_JSON_CHARS);
171                    Part::ToolCall(compacted_call)
172                }
173                Part::ToolResult(tool_result) => {
174                    let filtered = tool_result.filter_for_save();
175                    let compacted_tool_parts = filtered
176                        .parts
177                        .iter()
178                        .map(|tool_part| match tool_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                            // Keep artifact references; drop inline images from rolling context.
182                            Part::Image(_) => Part::Text(
183                                "[Image omitted from history; use artifact/reference if needed]"
184                                    .to_string(),
185                            ),
186                            other => other.clone(),
187                        })
188                        .collect();
189
190                    Part::ToolResult(ToolResponse {
191                        tool_call_id: filtered.tool_call_id,
192                        tool_name: filtered.tool_name,
193                        parts: compacted_tool_parts,
194                        parts_metadata: None,
195                    })
196                }
197                Part::Image(_) => {
198                    Part::Text("[Image omitted from history to reduce context size]".to_string())
199                }
200                Part::Artifact(artifact) => Part::Artifact(artifact.clone()),
201            })
202            .collect();
203
204        Self {
205            step_id: self.step_id.clone(),
206            parts: compacted_parts,
207            status: self.status.clone(),
208            reason: self.reason.as_ref().map(|r| truncate(r, MAX_TEXT_CHARS)),
209            timestamp: self.timestamp,
210        }
211    }
212}
213
214#[derive(Debug, Clone, JsonSchema, Serialize, Deserialize, PartialEq, Eq)]
215#[serde(rename_all = "snake_case")]
216pub enum ExecutionStatus {
217    Success,
218    Failed,
219    Rejected,
220    InputRequired,
221}
222
223impl From<ExecutionStatus> for TaskStatus {
224    fn from(val: ExecutionStatus) -> Self {
225        match val {
226            ExecutionStatus::Success => TaskStatus::Completed,
227            ExecutionStatus::Failed => TaskStatus::Failed,
228            ExecutionStatus::Rejected => TaskStatus::Canceled,
229            ExecutionStatus::InputRequired => TaskStatus::InputRequired,
230        }
231    }
232}
233
234pub enum ToolResultWithSkip {
235    ToolResult(ToolResponse),
236    // Skip tool call if it is external
237    Skip {
238        tool_call_id: String,
239        reason: String,
240    },
241}
242
243pub fn from_tool_results(tool_results: Vec<ToolResultWithSkip>) -> Vec<Part> {
244    tool_results
245        .iter()
246        .filter_map(|result| match result {
247            ToolResultWithSkip::ToolResult(tool_result) => {
248                // Simply extract parts from the tool response
249                Some(tool_result.parts.clone())
250            }
251            _ => None,
252        })
253        .flatten()
254        .collect()
255}
256
257#[derive(Debug, Clone, Serialize, Deserialize, Default)]
258pub struct ContextUsage {
259    pub tokens: u32,
260    pub input_tokens: u32,
261    pub output_tokens: u32,
262    /// Tokens read from provider cache (e.g., Anthropic prompt caching)
263    #[serde(default)]
264    pub cached_tokens: u32,
265    pub current_iteration: usize,
266    pub context_size: ContextSize,
267    /// Model used for LLM calls in this context
268    #[serde(default)]
269    pub model: Option<String>,
270}
271
272#[derive(Debug, Clone, Serialize, Deserialize, Default)]
273pub struct ContextSize {
274    pub message_count: usize,
275    pub message_chars: usize,
276    pub message_estimated_tokens: usize,
277    pub execution_history_count: usize,
278    pub execution_history_chars: usize,
279    pub execution_history_estimated_tokens: usize,
280    pub scratchpad_chars: usize,
281    pub scratchpad_estimated_tokens: usize,
282    pub total_chars: usize,
283    pub total_estimated_tokens: usize,
284    /// Per-agent context size breakdown
285    pub agent_breakdown: std::collections::HashMap<String, AgentContextSize>,
286}
287
288#[derive(Debug, Clone, Serialize, Deserialize, Default)]
289pub struct AgentContextSize {
290    pub agent_id: String,
291    pub task_count: usize,
292    pub execution_history_count: usize,
293    pub execution_history_chars: usize,
294    pub execution_history_estimated_tokens: usize,
295    pub scratchpad_chars: usize,
296    pub scratchpad_estimated_tokens: usize,
297}
298
299/// Enriched execution history entry that includes context metadata
300#[derive(Debug, Clone, Serialize, Deserialize)]
301pub struct ExecutionHistoryEntry {
302    pub thread_id: String, // Conversation context
303    pub task_id: String,   // Individual user task/request
304    pub run_id: String,    // Specific execution strand
305    pub execution_result: ExecutionResult,
306    pub stored_at: i64, // When this was stored
307}
308
309/// Entry for scratchpad formatting
310#[derive(Debug, Clone, Serialize, Deserialize)]
311pub struct ScratchpadEntry {
312    pub timestamp: i64,
313    #[serde(flatten)]
314    pub entry_type: ScratchpadEntryType,
315    pub task_id: String,
316    #[serde(default)]
317    pub parent_task_id: Option<String>,
318    pub entry_kind: Option<String>,
319}
320
321/// Type of scratchpad entry - only for Thought/Action/Observation tracking
322#[derive(Debug, Clone, Serialize, Deserialize)]
323#[serde(rename_all = "snake_case", tag = "type", content = "data")]
324pub enum ScratchpadEntryType {
325    #[serde(rename = "task")]
326    Task(Vec<Part>),
327    #[serde(rename = "plan")]
328    PlanStep(PlanStep),
329    #[serde(rename = "execution")]
330    Execution(ExecutionHistoryEntry),
331    /// Compressed summary produced by Tier 2 (semantic) compaction
332    #[serde(rename = "summary")]
333    Summary(CompactionSummary),
334}
335
336/// Summary produced by semantic compaction of older scratchpad entries
337#[derive(Debug, Clone, Serialize, Deserialize)]
338pub struct CompactionSummary {
339    /// LLM-generated summary of compacted history
340    pub summary_text: String,
341    /// Number of entries that were summarized
342    pub entries_summarized: usize,
343    /// Timestamp range of summarized entries
344    pub from_timestamp: i64,
345    pub to_timestamp: i64,
346    /// Token count saved by this compaction
347    pub tokens_saved: usize,
348}
349
350#[cfg(test)]
351mod tests {
352    use super::*;
353    use serde_json::json;
354
355    #[test]
356    fn test_scratchpad_large_observation_issue() {
357        println!("=== TESTING LARGE DATA OBSERVATION IN SCRATCHPAD ===");
358
359        // Create a very large tool response observation (similar to search results)
360        let large_data = json!({
361            "results": (0..100).map(|i| json!({
362                "id": i,
363                "name": format!("Minister {}", i),
364                "email": format!("minister{}@gov.sg", i),
365                "portfolio": format!("Ministry of Complex Affairs {}", i),
366                "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),
367            })).collect::<Vec<_>>()
368        });
369
370        println!(
371            "Large data size: {} bytes",
372            serde_json::to_string(&large_data).unwrap().len()
373        );
374
375        // Test 1: Direct Part::Data (BROKEN - causes scratchpad bloat)
376        let execution_result_data = ExecutionResult {
377            step_id: "test-step-1".to_string(),
378            parts: vec![Part::Data(large_data.clone())],
379            status: ExecutionStatus::Success,
380            reason: None,
381            timestamp: 1234567890,
382        };
383
384        let observation_data = execution_result_data.as_observation();
385        println!(
386            "🚨 BROKEN: Direct Part::Data observation size: {} chars",
387            observation_data.len()
388        );
389        println!(
390            "Preview (first 200 chars): {}",
391            &observation_data.chars().take(200).collect::<String>()
392        );
393
394        // Test 2: File metadata (GOOD - concise)
395        let file_metadata = crate::filesystem::FileMetadata {
396            file_id: "large-search-results.json".to_string(),
397            relative_path: "thread123/task456/large-search-results.json".to_string(),
398            size: serde_json::to_string(&large_data).unwrap().len() as u64,
399            content_type: Some("application/json".to_string()),
400            original_filename: Some("search_results.json".to_string()),
401            created_at: chrono::Utc::now(),
402            updated_at: chrono::Utc::now(),
403            checksum: Some("abc123".to_string()),
404            stats: None,
405            preview: Some("JSON search results with 100 minister entries".to_string()),
406        };
407
408        let execution_result_file = ExecutionResult {
409            step_id: "test-step-2".to_string(),
410            parts: vec![Part::Artifact(file_metadata)],
411            status: ExecutionStatus::Success,
412            reason: None,
413            timestamp: 1234567890,
414        };
415
416        let observation_file = execution_result_file.as_observation();
417        println!(
418            "✅ GOOD: File metadata observation size: {} chars",
419            observation_file.len()
420        );
421        println!("Content: {}", observation_file);
422
423        // Demonstrate the problem
424        println!("\n=== SCRATCHPAD IMPACT ===");
425        println!(
426            "❌ Direct approach adds {} chars to scratchpad (CAUSES LOOPS!)",
427            observation_data.len()
428        );
429        println!(
430            "✅ File metadata adds only {} chars to scratchpad",
431            observation_file.len()
432        );
433        println!(
434            "💡 Size reduction: {:.1}%",
435            (1.0 - (observation_file.len() as f64 / observation_data.len() as f64)) * 100.0
436        );
437
438        // This test shows the fix is working - observations are now truncated
439        assert!(observation_data.len() < 1000, "Large data is now truncated"); // Fixed expectation
440        assert!(
441            observation_file.len() < 300,
442            "File metadata stays reasonably concise"
443        ); // Updated for detailed format
444
445        println!("\n🚨 CONCLUSION: as_observation() needs to truncate large Part::Data!");
446    }
447
448    #[test]
449    fn test_observation_truncation_fix() {
450        println!("=== TESTING OBSERVATION TRUNCATION FIX ===");
451
452        // Test large data truncation
453        let large_data = json!({
454            "big_array": (0..200).map(|i| format!("item_{}", i)).collect::<Vec<_>>()
455        });
456
457        let execution_result = ExecutionResult {
458            step_id: "test-truncation".to_string(),
459            parts: vec![Part::Data(large_data)],
460            status: ExecutionStatus::Success,
461            reason: None,
462            timestamp: 1234567890,
463        };
464
465        let observation = execution_result.as_observation();
466        println!("Truncated observation size: {} chars", observation.len());
467        println!("Content: {}", observation);
468
469        // Should be truncated and include total char count
470        assert!(
471            observation.len() < 600,
472            "Observation should be truncated to <600 chars"
473        );
474        assert!(
475            observation.contains("truncated"),
476            "Should indicate truncation"
477        );
478        assert!(
479            observation.contains("total chars"),
480            "Should show total char count"
481        );
482
483        // Test long text truncation
484        let long_text = "This is a very long text. ".repeat(100);
485        let text_result = ExecutionResult {
486            step_id: "test-text-truncation".to_string(),
487            parts: vec![Part::Text(long_text.clone())],
488            status: ExecutionStatus::Success,
489            reason: None,
490            timestamp: 1234567890,
491        };
492
493        let text_observation = text_result.as_observation();
494        println!("Text observation size: {} chars", text_observation.len());
495        assert!(
496            text_observation.len() < 1100,
497            "Text should be truncated to ~1000 chars"
498        );
499        if long_text.len() > 1000 {
500            assert!(
501                text_observation.contains("truncated"),
502                "Long text should be truncated"
503            );
504        }
505
506        println!("✅ Observation truncation is working!");
507    }
508
509    #[test]
510    fn test_compact_for_history_filters_save_false_and_truncates_large_parts() {
511        let mut parts_metadata = std::collections::HashMap::new();
512        parts_metadata.insert(1, crate::PartMetadata { save: false });
513
514        let tool_response = ToolResponse {
515            tool_call_id: "call-1".to_string(),
516            tool_name: "search".to_string(),
517            parts: vec![
518                Part::Data(json!({"small": "kept"})),
519                Part::Data(json!({"secret": "do not persist"})),
520            ],
521            parts_metadata: Some(parts_metadata),
522        };
523
524        let huge = "x".repeat(6_000);
525        let execution_result = ExecutionResult {
526            step_id: "step-1".to_string(),
527            parts: vec![
528                Part::Text("y".repeat(2_500)),
529                Part::Data(json!({"huge": huge})),
530                Part::ToolResult(tool_response),
531            ],
532            status: ExecutionStatus::Success,
533            reason: Some("z".repeat(2_500)),
534            timestamp: 0,
535        };
536
537        let compacted = execution_result.compact_for_history();
538
539        assert_eq!(compacted.parts.len(), 3);
540        let text = match &compacted.parts[0] {
541            Part::Text(value) => value,
542            other => panic!("unexpected part: {:?}", other),
543        };
544        assert!(text.contains("[truncated"));
545
546        let data = match &compacted.parts[1] {
547            Part::Data(value) => value,
548            other => panic!("unexpected part: {:?}", other),
549        };
550        assert_eq!(data["truncated"], json!(true));
551
552        let tool = match &compacted.parts[2] {
553            Part::ToolResult(value) => value,
554            other => panic!("unexpected part: {:?}", other),
555        };
556        // save:false part should be removed.
557        assert_eq!(tool.parts.len(), 1);
558        assert!(tool.parts_metadata.is_none());
559    }
560}
distri_types/execution.rs

distri_types/
execution.rs
