Bound agent context and preprocess oversized image attachments

crates/galaxy_agent_core/src/types.rs

@@ -4,6 +4,9 @@ use serde::{Deserialize, Serialize};
 use serde_json::Value as JsonValue;
 
 pub const MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST: usize = 64_000;
+const MIN_PROVIDER_OUTPUT_DISCLOSURE_CHARS: usize = 27;
+const ESTIMATED_CHARS_PER_TOKEN: usize = 3;
+const MAX_SINGLE_OUTPUT_CONTEXT_FRACTION: usize = 4;
 
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
 pub struct ConversationMessage {
@@ -13,8 +16,51 @@ pub struct ConversationMessage {
 
 impl ConversationMessage {
     pub fn truncate_tool_results_for_provider_request(&mut self) {
-        truncate_tool_results_in_content(&mut self.content);
+        self.truncate_tool_results_for_provider_request_with_context(None);
     }
+
+    pub fn truncate_tool_results_for_provider_request_with_context(
+        &mut self,
+        max_context_tokens: Option<u32>,
+    ) {
+        truncate_tool_results_in_content(
+            &mut self.content,
+            provider_output_char_limit(max_context_tokens),
+        );
+    }
+
+    pub fn truncate_tagged_terminal_output_for_provider_request_with_context(
+        &mut self,
+        max_context_tokens: Option<u32>,
+    ) {
+        truncate_tagged_terminal_output_in_content(&mut self.content, max_context_tokens);
+    }
+}
+
+pub fn provider_output_char_limit(max_context_tokens: Option<u32>) -> usize {
+    let context_aware_limit = max_context_tokens
+        .and_then(|tokens| usize::try_from(tokens).ok())
+        .map(|tokens| {
+            tokens
+                .saturating_mul(ESTIMATED_CHARS_PER_TOKEN)
+                .saturating_div(MAX_SINGLE_OUTPUT_CONTEXT_FRACTION)
+        })
+        .unwrap_or(MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST);
+    context_aware_limit.clamp(
+        MIN_PROVIDER_OUTPUT_DISCLOSURE_CHARS,
+        MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST,
+    )
+}
+
+pub fn bounded_terminal_output_for_provider_request(
+    output: &str,
+    max_context_tokens: Option<u32>,
+) -> String {
+    truncate_provider_text(
+        output,
+        provider_output_char_limit(max_context_tokens),
+        "terminal output truncated for model context; full output remains available in the terminal",
+    )
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
@@ -323,15 +369,77 @@ pub enum AgentEvent {
     },
 }
 
-fn truncate_tool_results_in_content(content: &mut MessageContent) {
+fn truncate_tagged_terminal_output_in_content(
+    content: &mut MessageContent,
+    max_context_tokens: Option<u32>,
+) {
+    match content {
+        MessageContent::Text(text) => {
+            *text = truncate_tagged_terminal_output(text, max_context_tokens);
+        }
+        MessageContent::MultiPart(parts) => {
+            for part in parts {
+                if let ContentPart::Text(text) = part {
+                    *text = truncate_tagged_terminal_output(text, max_context_tokens);
+                }
+            }
+        }
+        MessageContent::ToolUse { .. } | MessageContent::ToolResult { .. } => {}
+    }
+}
+
+fn truncate_tagged_terminal_output(content: &str, max_context_tokens: Option<u32>) -> String {
+    const TAGS: [&str; 2] = ["[Terminal output:\n", "[Final terminal output:\n"];
+    const END_TAG: &str = "\n]\n";
+
+    let mut result = String::with_capacity(content.len());
+    let mut remaining = content;
+    loop {
+        let Some((tag_offset, tag)) = TAGS
+            .iter()
+            .filter_map(|tag| remaining.find(tag).map(|offset| (offset, *tag)))
+            .min_by_key(|(offset, _)| *offset)
+        else {
+            result.push_str(remaining);
+            break;
+        };
+        let output_start = tag_offset + tag.len();
+        result.push_str(&remaining[..output_start]);
+        let after_tag = &remaining[output_start..];
+        let mut output_ends = after_tag.match_indices(END_TAG);
+        let Some((output_end, _)) = output_ends.next() else {
+            result.push_str(&bounded_terminal_output_for_provider_request(
+                after_tag,
+                max_context_tokens,
+            ));
+            break;
+        };
+        if output_ends.next().is_some() {
+            return truncate_provider_text(
+                content,
+                provider_output_char_limit(max_context_tokens),
+                "message text truncated for model context; full content remains in conversation history",
+            );
+        }
+        result.push_str(&bounded_terminal_output_for_provider_request(
+            &after_tag[..output_end],
+            max_context_tokens,
+        ));
+        result.push_str(END_TAG);
+        remaining = &after_tag[output_end + END_TAG.len()..];
+    }
+    result
+}
+
+fn truncate_tool_results_in_content(content: &mut MessageContent, max_chars: usize) {
     match content {
         MessageContent::Text(_) | MessageContent::ToolUse { .. } => {}
-        MessageContent::ToolResult { content, .. } => truncate_tool_result_text(content),
+        MessageContent::ToolResult { content, .. } => truncate_tool_result_text(content, max_chars),
         MessageContent::MultiPart(parts) => {
             for part in parts {
                 match part {
                     ContentPart::ToolResult { content, .. } => {
-                        truncate_tool_result_text(content);
+                        truncate_tool_result_text(content, max_chars);
                     }
                     ContentPart::Text(_)
                     | ContentPart::Reasoning { .. }
@@ -343,16 +451,31 @@ fn truncate_tool_results_in_content(content: &mut MessageContent) {
     }
 }
 
-fn truncate_tool_result_text(content: &mut String) {
+fn truncate_tool_result_text(content: &mut String, max_chars: usize) {
+    *content = truncate_provider_text(content, max_chars, "tool result truncated");
+}
+
+fn truncate_provider_text(content: &str, max_chars: usize, reason: &str) -> String {
     let char_count = content.chars().count();
-    if char_count <= MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST {
-        return;
+    if char_count <= max_chars {
+        return content.to_string();
     }
 
-    let omitted_chars = char_count.saturating_sub(MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST);
-    let marker = format!("\n... [tool result truncated; omitted {omitted_chars} chars] ...\n");
+    let omitted_chars = char_count.saturating_sub(max_chars);
+    let marker = format!("\n... [{reason}; omitted {omitted_chars} chars] ...\n");
     let marker_chars = marker.chars().count();
-    let retained_chars = MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST.saturating_sub(marker_chars);
+    if max_chars <= marker_chars {
+        let short_reason = if reason.starts_with("terminal output") {
+            "terminal output truncated"
+        } else if reason.starts_with("message text") {
+            "message text truncated"
+        } else {
+            "tool result truncated"
+        };
+        let short_marker = format!("[{short_reason}]");
+        return short_marker.chars().take(max_chars).collect();
+    }
+    let retained_chars = max_chars.saturating_sub(marker_chars);
     let head_chars = retained_chars / 2;
     let tail_chars = retained_chars.saturating_sub(head_chars);
     let head: String = content.chars().take(head_chars).collect();
@@ -364,7 +487,7 @@ fn truncate_tool_result_text(content: &mut String) {
         .chars()
         .rev()
         .collect();
-    *content = format!("{head}{marker}{tail}");
+    format!("{head}{marker}{tail}")
 }
 
 #[cfg(test)]

crates/galaxy_agent_core/src/types_tests.rs

@@ -19,12 +19,116 @@ fn truncates_large_tool_results_for_provider_request() {
     let MessageContent::ToolResult { content, .. } = message.content else {
         panic!("expected tool result");
     };
-    assert!(content.len() <= MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST + 128);
+    assert!(content.chars().count() <= MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST);
     assert!(content.starts_with(prefix));
     assert!(content.ends_with(suffix));
     assert!(content.contains("tool result truncated"));
 }
 
+#[test]
+fn context_aware_tool_result_limit_preserves_the_original_message() {
+    let output = format!("start:{}:end", "x".repeat(20_000));
+    let mut message = ConversationMessage {
+        role: MessageRole::User,
+        content: MessageContent::ToolResult {
+            tool_use_id: "toolu_1".to_string(),
+            content: output.clone(),
+            is_error: false,
+        },
+    };
+
+    message.truncate_tool_results_for_provider_request_with_context(Some(8_000));
+
+    let MessageContent::ToolResult { content, .. } = message.content else {
+        panic!("expected tool result");
+    };
+    assert!(content.chars().count() <= provider_output_char_limit(Some(8_000)));
+    assert!(content.contains("tool result truncated"));
+    assert!(content.starts_with("start:"));
+    assert!(content.ends_with(":end"));
+    assert_eq!(output.chars().count(), 20_010);
+}
+
+#[test]
+fn terminal_output_bound_discloses_that_full_output_remains_in_terminal() {
+    let output = format!("first-line\n{}\nlast-line", "x".repeat(20_000));
+
+    let bounded = bounded_terminal_output_for_provider_request(&output, Some(8_000));
+
+    assert!(bounded.chars().count() <= provider_output_char_limit(Some(8_000)));
+    assert!(bounded.contains("terminal output truncated for model context"));
+    assert!(bounded.contains("full output remains available in the terminal"));
+    assert!(bounded.starts_with("first-line"));
+    assert!(bounded.ends_with("last-line"));
+    assert_eq!(output.chars().count(), 20_021);
+}
+
+#[test]
+fn terminal_output_bound_respects_character_budget_for_multibyte_text() {
+    let output = format!("first:{}:last", "界".repeat(20_000));
+
+    let bounded = bounded_terminal_output_for_provider_request(&output, Some(8_000));
+
+    assert!(bounded.chars().count() <= provider_output_char_limit(Some(8_000)));
+    assert!(bounded.starts_with("first:"));
+    assert!(bounded.ends_with(":last"));
+    assert!(bounded.contains("terminal output truncated for model context"));
+}
+
+#[test]
+fn tagged_terminal_output_is_bounded_without_mutating_surrounding_text() {
+    let output = format!("first:{}:last", "x".repeat(20_000));
+    let mut message = ConversationMessage {
+        role: MessageRole::User,
+        content: MessageContent::Text(format!(
+            "[Running command: verbose]\n[Terminal output:\n{output}\n]\nContinue the task"
+        )),
+    };
+
+    message.truncate_tagged_terminal_output_for_provider_request_with_context(Some(8_000));
+    let once = message.clone();
+    message.truncate_tagged_terminal_output_for_provider_request_with_context(Some(8_000));
+
+    assert_eq!(message, once);
+    let MessageContent::Text(content) = message.content else {
+        panic!("expected text message");
+    };
+    assert!(content.starts_with("[Running command: verbose]"));
+    assert!(content.ends_with("Continue the task"));
+    assert!(content.contains("terminal output truncated for model context"));
+    assert!(content.contains("first:"));
+    assert!(content.contains(":last"));
+    assert!(!content.contains(&output));
+}
+
+#[test]
+fn ambiguous_tagged_terminal_output_bounds_the_whole_provider_message() {
+    let output = format!("first\n]\n{}\n]\nlast", "x".repeat(20_000));
+    let mut message = ConversationMessage {
+        role: MessageRole::User,
+        content: MessageContent::Text(format!("[Terminal output:\n{output}\n]\nContinue the task")),
+    };
+
+    message.truncate_tagged_terminal_output_for_provider_request_with_context(Some(8_000));
+
+    let MessageContent::Text(content) = message.content else {
+        panic!("expected text message");
+    };
+    assert!(content.chars().count() <= provider_output_char_limit(Some(8_000)));
+    assert!(content.contains("message text truncated for model context"));
+    assert!(content.contains("full content remains in conversation history"));
+    assert!(!content.contains(&output));
+}
+
+#[test]
+fn terminal_output_bound_keeps_a_disclosure_for_tiny_budgets() {
+    let bounded = bounded_terminal_output_for_provider_request(&"x".repeat(200), Some(4));
+
+    assert_eq!(bounded, "[terminal output truncated]");
+    assert_eq!(bounded.chars().count(), provider_output_char_limit(Some(4)));
+    assert!(!bounded.contains(&"x".repeat(200)));
+}
+
 #[test]
 fn usage_total_combines_disjoint_cache_and_miss_buckets() {
     let usage = Usage {