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 const MAX_DATA_CHARS: usize = 500;
43 const MAX_TEXT_CHARS: usize = 1000;
44
45 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 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 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 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 pub const MAX_TOOL_RESULT_TOKENS: usize = 500;
228
229 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 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 {
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 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 #[serde(default)]
293 pub cached_tokens: u32,
294 pub current_iteration: usize,
295 pub context_size: ContextSize,
296 #[serde(default)]
298 pub model: Option<String>,
299 #[serde(default)]
301 pub context_budget: ContextBudget,
302 #[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#[derive(Debug, Clone, Serialize, Deserialize, Default)]
320pub struct ContextBudget {
321 pub system_prompt_static_tokens: usize,
323 pub system_prompt_dynamic_tokens: usize,
325 pub tool_schema_tokens: usize,
327 pub deferred_tool_tokens: usize,
329 pub skill_listing_tokens: usize,
331 pub conversation_tokens: usize,
333 pub tool_result_tokens: usize,
335 pub context_window_size: usize,
337 pub static_prefix_cache_hit: bool,
339 #[serde(default)]
341 pub static_prefix_hash: Option<String>,
342}
343
344impl ContextBudget {
345 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 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 pub fn remaining_tokens(&self) -> usize {
366 self.context_window_size.saturating_sub(self.total_tokens())
367 }
368
369 pub fn is_warning(&self) -> bool {
371 self.utilization() > 0.80
372 }
373
374 pub fn is_critical(&self) -> bool {
376 self.utilization() > 0.90
377 }
378
379 pub fn deferred_savings(&self) -> usize {
381 0 }
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 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#[derive(Debug, Clone, Serialize, Deserialize)]
415pub struct ExecutionHistoryEntry {
416 pub thread_id: String, pub task_id: String, pub run_id: String, pub execution_result: ExecutionResult,
420 pub stored_at: i64, }
422
423#[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#[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 #[serde(rename = "summary")]
447 Summary(CompactionSummary),
448 #[serde(rename = "skill_context")]
450 SkillContext(SkillContextEntry),
451}
452
453#[derive(Debug, Clone, Serialize, Deserialize)]
455pub struct SkillContextEntry {
456 pub skill_id: String,
458 pub content: String,
460 pub reinjected_at: i64,
462}
463
464#[derive(Debug, Clone, Serialize, Deserialize)]
466pub struct CompactionSummary {
467 pub summary_text: String,
469 pub entries_summarized: usize,
471 pub from_timestamp: i64,
473 pub to_timestamp: i64,
474 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 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 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 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 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 assert!(observation_data.len() < 1000, "Large data is now truncated"); assert!(
569 observation_file.len() < 300,
570 "File metadata stays reasonably concise"
571 ); 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 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 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 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 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}