1use schemars::JsonSchema;
2use serde::{Deserialize, Serialize};
3use serde_json::json;
4
5use crate::{Part, PlanStep, TaskStatus, ToolResponse, core::FileType};
6
7#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
9pub enum ExecutionType {
10 Interleaved,
11 Retriable,
12 React,
13 Code,
14}
15
16#[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>, 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 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 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 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 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 pub const MAX_TOOL_RESULT_TOKENS: usize = 500;
249
250 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 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 {
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 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 #[serde(default)]
318 pub cached_tokens: u32,
319 pub current_iteration: usize,
320 pub context_size: ContextSize,
321 #[serde(default)]
323 pub model: Option<String>,
324 #[serde(default)]
326 pub context_budget: ContextBudget,
327 #[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#[derive(Debug, Clone, Serialize, Deserialize, Default)]
345pub struct ContextBudget {
346 pub system_prompt_static_tokens: usize,
348 pub system_prompt_dynamic_tokens: usize,
350 pub tool_schema_tokens: usize,
352 pub deferred_tool_tokens: usize,
354 pub skill_listing_tokens: usize,
356 pub conversation_tokens: usize,
358 pub tool_result_tokens: usize,
360 pub context_window_size: usize,
362 pub static_prefix_cache_hit: bool,
364 #[serde(default)]
366 pub static_prefix_hash: Option<String>,
367}
368
369impl ContextBudget {
370 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 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 pub fn remaining_tokens(&self) -> usize {
391 self.context_window_size.saturating_sub(self.total_tokens())
392 }
393
394 pub fn is_warning(&self) -> bool {
396 self.utilization() > 0.80
397 }
398
399 pub fn is_critical(&self) -> bool {
401 self.utilization() > 0.90
402 }
403
404 pub fn deferred_savings(&self) -> usize {
406 0 }
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 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#[derive(Debug, Clone, Serialize, Deserialize)]
440pub struct ExecutionHistoryEntry {
441 pub thread_id: String, pub task_id: String, pub run_id: String, pub execution_result: ExecutionResult,
445 pub stored_at: i64, }
447
448#[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#[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 #[serde(rename = "summary")]
472 Summary(CompactionSummary),
473 #[serde(rename = "skill_context")]
475 SkillContext(SkillContextEntry),
476}
477
478#[derive(Debug, Clone, Serialize, Deserialize)]
480pub struct SkillContextEntry {
481 pub skill_id: String,
483 pub content: String,
485 pub reinjected_at: i64,
487}
488
489#[derive(Debug, Clone, Serialize, Deserialize)]
491pub struct CompactionSummary {
492 pub summary_text: String,
494 pub entries_summarized: usize,
496 pub from_timestamp: i64,
498 pub to_timestamp: i64,
499 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 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 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 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 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 assert!(observation_data.len() < 1000, "Large data is now truncated"); assert!(
594 observation_file.len() < 300,
595 "File metadata stays reasonably concise"
596 ); 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 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 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 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 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}