From 1e284c625e67be96287b8caedf64819479210acd Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Mon, 24 Aug 2026 13:56:16 -0400 Subject: [PATCH] Simplify bounded tool output capture Apply cleanups from a reuse/simplification/efficiency review of the bounded-tool-output changes: - Share one MAX_RUN_EVENT_BODY_BYTES constant in fabro-types; the server body limit, the agent's serialized-output reservation, and the event headroom test all derive from it. - Rework truncation.rs around one split_head_tail helper: drop the hand-rolled ceil_char_boundary (std's is stable), the duplicate truncate_plain_output splitter and its dead Tail arm, and the head_bytes field with its sentinel values. - Return Cow from preview_tool_output and take retain_tool_output's input by value, so untruncated output crosses the pipeline without full copies. Measure serialized JSON size with a counting writer instead of materializing the payload. - Reuse fabro-llm's byte-token estimate (now public) instead of a third copy of the 4-bytes-per-token heuristic. - Take retain_tool_result's ToolResult by value and mutate content in place; extract the triplicated error retain-emit-truncate block into finish_error_result. - Share the shell retain-and-record sequence between the native and kimi shell tools as retain_shell_output. - Move OutputCaptureBuffer::into_parts to reuse the head allocation, skip the buffer round-trip in replay_exec_result when output fits, and replace daytona's byte-iterator suffix matching with contiguous slice comparisons behind one retained_slices accessor. - Make SessionBoundEmitter's fields private. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01TK3QTWQHiXhRbFwTr57LzX --- .../fabro-server/src/server/handler/events.rs | 3 +- lib/components/fabro-agent/src/event.rs | 6 +- .../fabro-agent/src/profiles/kimi_tools.rs | 11 +- .../fabro-agent/src/tool_execution.rs | 74 +++-- lib/components/fabro-agent/src/tools.rs | 30 ++- lib/components/fabro-agent/src/truncation.rs | 252 +++++++++--------- lib/components/fabro-llm/src/token_count.rs | 12 +- .../fabro-sandbox/src/daytona/mod.rs | 64 +++-- lib/components/fabro-sandbox/src/sandbox.rs | 61 +++-- .../fabro-types/src/run_event/mod.rs | 8 + 10 files changed, 266 insertions(+), 255 deletions(-) diff --git a/lib/apps/fabro-server/src/server/handler/events.rs b/lib/apps/fabro-server/src/server/handler/events.rs index 247823cb7..01e8085eb 100644 --- a/lib/apps/fabro-server/src/server/handler/events.rs +++ b/lib/apps/fabro-server/src/server/handler/events.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use axum::extract::DefaultBodyLimit; +use fabro_types::run_event::MAX_RUN_EVENT_BODY_BYTES; use fabro_types::{ RunEventDetailContent, RunEventDetailContentKind, RunEventDetailEnvelope, RunEventDetailResponse, @@ -16,8 +17,6 @@ use super::super::{ reject_if_archived, update_live_run_from_event, }; -const MAX_RUN_EVENT_BODY_BYTES: usize = 3 * 1024 * 1024; - pub(super) fn routes() -> Router> { Router::new() .route("/attach", get(attach_events)) diff --git a/lib/components/fabro-agent/src/event.rs b/lib/components/fabro-agent/src/event.rs index eaea439bd..e16f71e50 100644 --- a/lib/components/fabro-agent/src/event.rs +++ b/lib/components/fabro-agent/src/event.rs @@ -64,9 +64,9 @@ impl Default for Emitter { /// when a subagent's events are forwarded through its parent. #[derive(Clone)] pub struct SessionBoundEmitter { - pub emitter: Emitter, - pub session_id: String, - pub tool_call_id: Option, + emitter: Emitter, + session_id: String, + tool_call_id: Option, tool_output_stats: Arc>>, } diff --git a/lib/components/fabro-agent/src/profiles/kimi_tools.rs b/lib/components/fabro-agent/src/profiles/kimi_tools.rs index 6f13f2b57..07b6ec05e 100644 --- a/lib/components/fabro-agent/src/profiles/kimi_tools.rs +++ b/lib/components/fabro-agent/src/profiles/kimi_tools.rs @@ -31,9 +31,8 @@ use crate::sandbox::{GrepOptions, format_lines_numbered}; use crate::tool_registry::{RegisteredTool, ToolSource}; use crate::tools::{ DEFAULT_READ_LINES, emit_shell_process_completed, execute_grep, execute_shell_command, - grep_result_path, make_edit_file_tool, optional_usize_arg, required_str, + grep_result_path, make_edit_file_tool, optional_usize_arg, required_str, retain_shell_output, }; -use crate::truncation::{MAX_RETAINED_TOOL_OUTPUT_BYTES, retain_tool_output}; const DEFAULT_GREP_RESULTS: usize = 250; const MAX_GREP_RESULTS: usize = 2000; @@ -144,13 +143,7 @@ explicitly asked. Never run commands requiring superuser privileges unless expli let _ = write!(out, "Command failed with exit code: {code}"); } let is_success = result.is_success(); - let retained = retain_tool_output( - &out, - MAX_RETAINED_TOOL_OUTPUT_BYTES, - streaming.output_capture().omitted_bytes, - ); - ctx.record_tool_output_stats(retained.stats); - let out = retained.output; + let out = retain_shell_output(&ctx, &streaming, out); emit_shell_process_completed(&ctx, streaming).await; if is_success { Ok(out) } else { Err(out) } }) diff --git a/lib/components/fabro-agent/src/tool_execution.rs b/lib/components/fabro-agent/src/tool_execution.rs index 7b377e586..cc2943aaf 100644 --- a/lib/components/fabro-agent/src/tool_execution.rs +++ b/lib/components/fabro-agent/src/tool_execution.rs @@ -1,3 +1,4 @@ +use std::borrow::Cow; use std::sync::Arc; use fabro_llm::types::{ToolCall, ToolResult}; @@ -12,7 +13,8 @@ use crate::sandbox::{OutputCaptureStats, Sandbox}; use crate::session::ToolEnvProvider; use crate::tool_registry::{AgentEventEmitter, RegisteredTool, ToolContext, ToolRegistry}; use crate::truncation::{ - MAX_RETAINED_TOOL_OUTPUT_BYTES, preview_tool_output, truncate_tool_output, + MAX_RETAINED_TOOL_OUTPUT_BYTES, preview_tool_output, serialized_json_bytes, + truncate_tool_output, }; use crate::types::AgentEvent; @@ -267,7 +269,19 @@ fn error_tool_result_with_events( message: &str, ) -> ToolResult { emit_tool_call_started(emitter, session_id, tc); - let retained = retain_tool_result(&ToolResult::error(&tc.id, message), None); + finish_error_result(tc, emitter, session_id, config, message) +} + +/// Bound, emit, and truncate an error result for a tool call whose +/// started event was already emitted. +fn finish_error_result( + tc: &ToolCall, + emitter: &Emitter, + session_id: &str, + config: &SessionOptions, + message: &str, +) -> ToolResult { + let retained = retain_tool_result(ToolResult::error(&tc.id, message), None); emit_tool_call_result( emitter, session_id, @@ -403,15 +417,7 @@ async fn execute_and_emit_one_tool_with_lookup( emit_tool_call_started(emitter, session_id, tc); if let Some(reason) = access_denial { - let retained = retain_tool_result(&ToolResult::error(&tc.id, &reason), None); - emit_tool_call_result( - emitter, - session_id, - tc, - &retained.result, - retained.output_stats, - ); - return truncate_tool_result(&retained.result, &tc.name, config); + return finish_error_result(tc, emitter, session_id, config, &reason); } // Pre-tool-use hook @@ -423,15 +429,7 @@ async fn execute_and_emit_one_tool_with_lookup( debug!(tool = %tc.name, hook_event = "pre_tool_use", ?decision, duration_ms = elapsed, "Tool hook complete"); if let ToolHookDecision::Block { reason } = decision { - let retained = retain_tool_result(&ToolResult::error(&tc.id, &reason), None); - emit_tool_call_result( - emitter, - session_id, - tc, - &retained.result, - retained.output_stats, - ); - return truncate_tool_result(&retained.result, &tc.name, config); + return finish_error_result(tc, emitter, session_id, config, &reason); } } @@ -447,7 +445,7 @@ async fn execute_and_emit_one_tool_with_lookup( agent_tool_runtime, ) .await; - let retained = retain_tool_result(&executed.result, executed.output_stats); + let retained = retain_tool_result(executed.result, executed.output_stats); let result = retained.result; emit_tool_call_result(emitter, session_id, tc, &result, retained.output_stats); @@ -484,32 +482,25 @@ struct RetainedToolResult { /// Bound model-native tool output before it reaches hooks, events, or history. fn retain_tool_result( - result: &ToolResult, + mut result: ToolResult, previous_stats: Option, ) -> RetainedToolResult { - let (retained_content, output_stats) = match &result.content { + let output_stats = match &mut result.content { serde_json::Value::String(output) => { let previously_omitted = previous_stats.map_or(0, |stats| stats.omitted_bytes); - let retained = + let previewed = preview_tool_output(output, MAX_RETAINED_TOOL_OUTPUT_BYTES, previously_omitted); - (serde_json::Value::String(retained.output), retained.stats) - } - other => { - let byte_count = serde_json::to_vec(other) - .expect("serde_json::Value always serializes") - .len(); - (other.clone(), OutputCaptureStats::complete(byte_count)) + let stats = previewed.stats; + if let Cow::Owned(previewed_output) = previewed.output { + *output = previewed_output; + } + stats } + other => OutputCaptureStats::complete(serialized_json_bytes(other)), }; RetainedToolResult { - result: ToolResult { - tool_call_id: result.tool_call_id.clone(), - content: retained_content, - is_error: result.is_error, - image_data: result.image_data.clone(), - image_media_type: result.image_media_type.clone(), - }, + result, output_stats, } } @@ -645,7 +636,7 @@ mod tests { use async_trait::async_trait; use fabro_llm::types::{ToolCall, ToolDefinition}; use fabro_model::AgentProfileKind; - use fabro_types::run_event::AgentToolCompletedProps; + use fabro_types::run_event::{AgentToolCompletedProps, MAX_RUN_EVENT_BODY_BYTES}; use tokio::sync::broadcast; use super::*; @@ -1114,8 +1105,11 @@ mod tests { let serialized_event_bytes = serde_json::to_vec(&run_event) .expect("run event serializes") .len(); + // Leave at least 1 MiB of envelope headroom under the server's + // run-event body limit. + let event_body_budget = MAX_RUN_EVENT_BODY_BYTES - 1024 * 1024; assert!( - serialized_event_bytes < 2 * 1024 * 1024, + serialized_event_bytes < event_body_budget, "serialized event was {serialized_event_bytes} bytes" ); } diff --git a/lib/components/fabro-agent/src/tools.rs b/lib/components/fabro-agent/src/tools.rs index 24ab550d2..24ad2c80f 100644 --- a/lib/components/fabro-agent/src/tools.rs +++ b/lib/components/fabro-agent/src/tools.rs @@ -13,7 +13,7 @@ use tokio::task; use crate::config::NativeToolOptions; use crate::sandbox::{ExecStreamingResult, GrepOptions}; use crate::tool_registry::{RegisteredTool, ToolContext, ToolRegistry, ToolSource}; -use crate::truncation::{MAX_RETAINED_TOOL_OUTPUT_BYTES, RetainedToolOutput, retain_tool_output}; +use crate::truncation::{MAX_RETAINED_TOOL_OUTPUT_BYTES, retain_tool_output}; use crate::types::AgentEvent; use crate::web_search::{SearchBackend, make_web_search_tool}; @@ -320,15 +320,29 @@ pub(crate) async fn run_shell_command( cwd: Option<&str>, ) -> Result { let streaming = execute_shell_command(ctx, command, timeout_ms, cwd).await?; - let retained = render_shell_result(&streaming); - ctx.record_tool_output_stats(retained.stats); - let text = retained.output; + let text = retain_shell_output(ctx, &streaming, render_shell_result(&streaming)); let is_success = streaming.result.is_success(); emit_shell_process_completed(ctx, streaming).await; if is_success { Ok(text) } else { Err(text) } } +/// Bound rendered shell output to the retention budget and record the capture +/// stats for the executing tool call. +pub(crate) fn retain_shell_output( + ctx: &ToolContext, + streaming: &ExecStreamingResult, + output: String, +) -> String { + let retained = retain_tool_output( + output, + MAX_RETAINED_TOOL_OUTPUT_BYTES, + streaming.output_capture().omitted_bytes, + ); + ctx.record_tool_output_stats(retained.stats); + retained.output +} + /// Emit the subordinate process outcome after model-facing output has been /// rendered. Consumes the raw result so redaction does not require cloning /// potentially large process output. @@ -372,7 +386,7 @@ pub(crate) async fn emit_shell_process_completed( /// Renders the model-facing shell result: termination, exit code, duration, /// and provider-honest output sections. Metadata stays at the head and /// `stderr` at the tail so head/tail truncation preserves both. -fn render_shell_result(streaming: &ExecStreamingResult) -> RetainedToolOutput { +fn render_shell_result(streaming: &ExecStreamingResult) -> String { let result = &streaming.result; let mut output = format!( "Termination: {}\nExit code: {}\nDuration: {}ms\n", @@ -392,11 +406,7 @@ fn render_shell_result(streaming: &ExecStreamingResult) -> RetainedToolOutput { } else if !result.stdout.is_empty() { let _ = write!(output, "output (combined):\n{}\n", result.stdout); } - retain_tool_output( - &output, - MAX_RETAINED_TOOL_OUTPUT_BYTES, - streaming.output_capture().omitted_bytes, - ) + output } #[must_use] diff --git a/lib/components/fabro-agent/src/truncation.rs b/lib/components/fabro-agent/src/truncation.rs index 6e155da86..719cda287 100644 --- a/lib/components/fabro-agent/src/truncation.rs +++ b/lib/components/fabro-agent/src/truncation.rs @@ -1,15 +1,43 @@ +use std::borrow::Cow; + +use fabro_llm::token_count; +use fabro_types::run_event::MAX_RUN_EVENT_BODY_BYTES; +use serde::Serialize; + use crate::config::SessionOptions; use crate::sandbox::OutputCaptureStats; use crate::tool_permissions::canonical_tool_name; pub(crate) const MAX_RETAINED_TOOL_OUTPUT_BYTES: usize = 1024 * 1024; -pub(crate) const MAX_SERIALIZED_TOOL_OUTPUT_BYTES: usize = 3 * 1024 * 1024 / 2; +/// Reserve half the run-event body limit for serialized tool output; the +/// other half is headroom for the rest of the event envelope. +pub(crate) const MAX_SERIALIZED_TOOL_OUTPUT_BYTES: usize = MAX_RUN_EVENT_BODY_BYTES / 2; -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug)] pub(crate) struct RetainedToolOutput { pub output: String, pub stats: OutputCaptureStats, - head_bytes: usize, +} + +/// Model-facing preview of a tool output. Borrows the input when no +/// truncation notice was needed. +#[derive(Debug)] +pub(crate) struct PreviewedToolOutput<'a> { + pub output: Cow<'a, str>, + pub stats: OutputCaptureStats, +} + +/// Boundaries of an equal-sized UTF-8 head and tail fitting `max_bytes`, or +/// `None` when `output` already fits. +fn split_head_tail(output: &str, max_bytes: usize) -> Option<(usize, usize)> { + if output.len() <= max_bytes { + return None; + } + let head_budget = max_bytes / 2; + let tail_budget = max_bytes - head_budget; + let head_end = output.floor_char_boundary(head_budget); + let tail_start = output.ceil_char_boundary(output.len() - tail_budget); + Some((head_end, tail_start)) } /// Keep an equal-sized UTF-8 prefix and suffix within a byte budget. @@ -18,44 +46,34 @@ pub(crate) struct RetainedToolOutput { /// discarded before the rendered result was assembled. #[must_use] pub(crate) fn retain_tool_output( - output: &str, + output: String, max_bytes: usize, previously_omitted_bytes: usize, ) -> RetainedToolOutput { let observed_bytes = output.len().saturating_add(previously_omitted_bytes); - if output.len() <= max_bytes { + let Some((head_end, tail_start)) = split_head_tail(&output, max_bytes) else { return RetainedToolOutput { - output: output.to_string(), - stats: OutputCaptureStats { + stats: OutputCaptureStats { observed_bytes, retained_bytes: output.len(), omitted_bytes: previously_omitted_bytes, }, - head_bytes: if previously_omitted_bytes == 0 { - output.len() - } else { - output.floor_char_boundary(output.len() / 2) - }, + output, }; - } + }; - let head_budget = max_bytes / 2; - let tail_budget = max_bytes.saturating_sub(head_budget); - let head_end = output.floor_char_boundary(head_budget); - let tail_start = ceil_char_boundary(output, output.len().saturating_sub(tail_budget)); - let retained_bytes = head_end.saturating_add(output.len().saturating_sub(tail_start)); + let retained_bytes = head_end + (output.len() - tail_start); let mut retained = String::with_capacity(retained_bytes); retained.push_str(&output[..head_end]); retained.push_str(&output[tail_start..]); RetainedToolOutput { - output: retained, - stats: OutputCaptureStats { + output: retained, + stats: OutputCaptureStats { observed_bytes, retained_bytes, omitted_bytes: observed_bytes.saturating_sub(retained_bytes), }, - head_bytes: head_end, } } @@ -66,30 +84,64 @@ pub(crate) fn preview_tool_output( output: &str, max_bytes: usize, previously_omitted_bytes: usize, -) -> RetainedToolOutput { +) -> PreviewedToolOutput<'_> { + let observed_bytes = output.len().saturating_add(previously_omitted_bytes); let mut content_budget = max_bytes; loop { - let retained = retain_tool_output(output, content_budget, previously_omitted_bytes); - let rendered = if retained.stats.omitted_bytes == 0 { - retained.output.clone() + let (head_end, tail_start, stats) = + if let Some((head_end, tail_start)) = split_head_tail(output, content_budget) { + let retained_bytes = head_end + (output.len() - tail_start); + (head_end, tail_start, OutputCaptureStats { + observed_bytes, + retained_bytes, + omitted_bytes: observed_bytes.saturating_sub(retained_bytes), + }) + } else { + // The whole output fits. A notice is still rendered when the + // stream itself omitted bytes; equal-sized retention keeps + // that omission gap at the midpoint. + let mid = output.floor_char_boundary(output.len() / 2); + (mid, mid, OutputCaptureStats { + observed_bytes, + retained_bytes: output.len(), + omitted_bytes: previously_omitted_bytes, + }) + }; + let rendered: Cow<'_, str> = if stats.omitted_bytes == 0 { + Cow::Borrowed(output) } else { - render_retained_output(&retained) + Cow::Owned(render_truncated_segments( + &output[..head_end], + &output[tail_start..], + stats, + None, + )) }; - let serialized_bytes = serialized_json_string_bytes(&rendered); + let serialized_bytes = serialized_json_bytes(rendered.as_ref()); if rendered.len() <= max_bytes && serialized_bytes <= MAX_SERIALIZED_TOOL_OUTPUT_BYTES { - if retained.stats.omitted_bytes == 0 { - return retained; - } - return RetainedToolOutput { - output: rendered, - stats: retained.stats, - head_bytes: 0, + return PreviewedToolOutput { + output: rendered, + stats, }; } - let mut next_budget = content_budget; + let Some(reduced_budget) = content_budget.checked_sub(1) else { + // The content budget is exhausted and the notice text alone still + // overflows. Hard-cut the rendered notice to fit. + let output = match split_head_tail(&rendered, max_bytes) { + Some((head_end, tail_start)) => { + format!("{}{}", &rendered[..head_end], &rendered[tail_start..]) + } + None => rendered.into_owned(), + }; + return PreviewedToolOutput { + output: Cow::Owned(output), + stats, + }; + }; + let mut next_budget = reduced_budget; if rendered.len() > max_bytes { - let excess = rendered.len().saturating_sub(max_bytes); + let excess = rendered.len() - max_bytes; next_budget = next_budget.min(content_budget.saturating_sub(excess)); } if serialized_bytes > MAX_SERIALIZED_TOOL_OUTPUT_BYTES { @@ -100,28 +152,27 @@ pub(crate) fn preview_tool_output( .unwrap_or(0); next_budget = next_budget.min(scaled_budget); } - next_budget = next_budget.min(content_budget.saturating_sub(1)); - if next_budget == content_budget { - return RetainedToolOutput { - output: truncate_plain_output(&rendered, max_bytes, TruncationMode::HeadTail), - stats: retained.stats, - head_bytes: 0, - }; - } content_budget = next_budget; } } -fn serialized_json_string_bytes(output: &str) -> usize { - serde_json::to_vec(output) - .expect("strings always serialize as JSON") - .len() -} +/// Serialized JSON size in bytes, counted without materializing the payload. +pub(crate) fn serialized_json_bytes(value: &T) -> usize { + struct CountingWriter(usize); + impl std::io::Write for CountingWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.0 += buf.len(); + Ok(buf.len()) + } -fn render_retained_output(retained: &RetainedToolOutput) -> String { - let head = &retained.output[..retained.head_bytes]; - let tail = &retained.output[retained.head_bytes..]; - render_truncated_segments(head, tail, retained.stats, None) + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + let mut writer = CountingWriter(0); + serde_json::to_writer(&mut writer, value).expect("JSON tool output always serializes"); + writer.0 } fn render_truncated_segments( @@ -130,8 +181,8 @@ fn render_truncated_segments( stats: OutputCaptureStats, line_count_omitted: Option, ) -> String { - let original_tokens = approximate_tokens(stats.observed_bytes); - let omitted_tokens = approximate_tokens(stats.omitted_bytes); + let original_tokens = token_count::estimate_byte_tokens(stats.observed_bytes); + let omitted_tokens = token_count::estimate_byte_tokens(stats.omitted_bytes); let middle_marker = line_count_omitted.map_or_else( || format!("... approximately {omitted_tokens} tokens truncated ..."), |lines| { @@ -146,18 +197,6 @@ fn render_truncated_segments( ) } -fn approximate_tokens(bytes: usize) -> usize { - bytes.div_ceil(4) -} - -fn ceil_char_boundary(output: &str, index: usize) -> usize { - let mut index = index.min(output.len()); - while index < output.len() && !output.is_char_boundary(index) { - index += 1; - } - index -} - #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum TruncationMode { HeadTail, @@ -193,63 +232,28 @@ fn default_truncation_mode(tool_name: &str) -> TruncationMode { #[must_use] pub fn truncate_output(output: &str, max_chars: usize, mode: TruncationMode) -> String { - if output.len() <= max_chars { + let Some((head_end, tail_start)) = split_head_tail(output, max_chars) else { return output.to_string(); - } + }; - match mode { - TruncationMode::HeadTail => { - let half = max_chars / 2; - let head_end = output.floor_char_boundary(half); - let tail_start = ceil_char_boundary(output, output.len().saturating_sub(half)); - let head = &output[..head_end]; - let tail = &output[tail_start..]; - let retained_bytes = head.len().saturating_add(tail.len()); - render_truncated_segments( - head, - tail, - OutputCaptureStats { - observed_bytes: output.len(), - retained_bytes, - omitted_bytes: output.len().saturating_sub(retained_bytes), - }, - None, - ) - } + let (head, tail) = match mode { + TruncationMode::HeadTail => (&output[..head_end], &output[tail_start..]), TruncationMode::Tail => { - let tail_start = ceil_char_boundary(output, output.len().saturating_sub(max_chars)); - let tail = &output[tail_start..]; - render_truncated_segments( - "", - tail, - OutputCaptureStats { - observed_bytes: output.len(), - retained_bytes: tail.len(), - omitted_bytes: output.len().saturating_sub(tail.len()), - }, - None, - ) + let tail_start = output.ceil_char_boundary(output.len() - max_chars); + ("", &output[tail_start..]) } - } -} - -fn truncate_plain_output(output: &str, max_bytes: usize, mode: TruncationMode) -> String { - if output.len() <= max_bytes { - return output.to_string(); - } - - match mode { - TruncationMode::HeadTail => { - let half = max_bytes / 2; - let head_end = output.floor_char_boundary(half); - let tail_start = ceil_char_boundary(output, output.len().saturating_sub(half)); - format!("{}{}", &output[..head_end], &output[tail_start..]) - } - TruncationMode::Tail => { - let tail_start = ceil_char_boundary(output, output.len().saturating_sub(max_bytes)); - output[tail_start..].to_string() - } - } + }; + let retained_bytes = head.len().saturating_add(tail.len()); + render_truncated_segments( + head, + tail, + OutputCaptureStats { + observed_bytes: output.len(), + retained_bytes, + omitted_bytes: output.len().saturating_sub(retained_bytes), + }, + None, + ) } #[must_use] @@ -316,7 +320,7 @@ mod tests { #[test] fn retained_tool_output_keeps_equal_head_and_tail() { - let retained = retain_tool_output("abcdefghijkl", 8, 0); + let retained = retain_tool_output("abcdefghijkl".to_string(), 8, 0); assert_eq!(retained.output, "abcdijkl"); assert_eq!(retained.stats.observed_bytes, 12); @@ -326,7 +330,7 @@ mod tests { #[test] fn retained_tool_output_stays_within_budget_at_utf8_boundaries() { - let retained = retain_tool_output("aa😀😀zz", 7, 3); + let retained = retain_tool_output("aa😀😀zz".to_string(), 7, 3); assert!(retained.output.len() <= 7, "{}", retained.output.len()); assert!(retained.output.starts_with("aa")); @@ -380,10 +384,10 @@ mod tests { "\0".repeat(MAX_RETAINED_TOOL_OUTPUT_BYTES - "HEADTAIL".len()) ); assert_eq!(output.len(), MAX_RETAINED_TOOL_OUTPUT_BYTES); - assert!(serialized_json_string_bytes(&output) > MAX_SERIALIZED_TOOL_OUTPUT_BYTES); + assert!(serialized_json_bytes(output.as_str()) > MAX_SERIALIZED_TOOL_OUTPUT_BYTES); let preview = preview_tool_output(&output, MAX_RETAINED_TOOL_OUTPUT_BYTES, 0); - let serialized_bytes = serialized_json_string_bytes(&preview.output); + let serialized_bytes = serialized_json_bytes(preview.output.as_ref()); assert!(preview.output.len() <= MAX_RETAINED_TOOL_OUTPUT_BYTES); assert!( diff --git a/lib/components/fabro-llm/src/token_count.rs b/lib/components/fabro-llm/src/token_count.rs index 58d9ec74d..fa64abda7 100644 --- a/lib/components/fabro-llm/src/token_count.rs +++ b/lib/components/fabro-llm/src/token_count.rs @@ -221,7 +221,7 @@ impl Estimator { let mut tokens = estimate_text_tokens(&result.tool_call_id) + estimate_json_tokens(&result.content); if let Some(image_data) = &result.image_data { - tokens += estimate_embedded_bytes(image_data.len()); + tokens += estimate_byte_tokens(image_data.len()); self.warn( MEDIA_ESTIMATE_WARNING, "Media content couldn't be precisely tokenized; total is approximate.", @@ -243,7 +243,7 @@ impl Estimator { + image .data .as_ref() - .map_or(2000, |data| estimate_embedded_bytes(data.len()).max(2000)) + .map_or(2000, |data| estimate_byte_tokens(data.len()).max(2000)) } fn estimate_audio(&mut self, audio: &AudioData) -> usize { @@ -252,7 +252,7 @@ impl Estimator { + audio .data .as_ref() - .map_or(2000, |data| estimate_embedded_bytes(data.len())) + .map_or(2000, |data| estimate_byte_tokens(data.len())) } fn estimate_document(&mut self, document: &DocumentData) -> usize { @@ -265,7 +265,7 @@ impl Estimator { + document .data .as_ref() - .map_or(2000, |data| estimate_embedded_bytes(data.len())) + .map_or(2000, |data| estimate_byte_tokens(data.len())) } fn estimate_media_common(&mut self, url: Option<&str>, media_type: Option<&str>) -> usize { @@ -292,7 +292,9 @@ fn estimate_tool(tool: &ToolDefinition) -> usize { + estimate_json_tokens(&tool.parameters) } -fn estimate_embedded_bytes(byte_len: usize) -> usize { +/// Approximate the token cost of a raw byte payload (4 bytes per token). +#[must_use] +pub fn estimate_byte_tokens(byte_len: usize) -> usize { byte_len.div_ceil(4) } diff --git a/lib/components/fabro-sandbox/src/daytona/mod.rs b/lib/components/fabro-sandbox/src/daytona/mod.rs index 6da9c6e72..a7a0a8d12 100644 --- a/lib/components/fabro-sandbox/src/daytona/mod.rs +++ b/lib/components/fabro-sandbox/src/daytona/mod.rs @@ -33,8 +33,8 @@ use crate::push_credentials::{self, PushCredentialState}; use crate::redact::redact_auth_url; use crate::sandbox::{ self, BASH_ENV_VAR, BASH_PROBE_MARKER, BASH_PROBE_SCRIPT, BASH_PROBE_TIMEOUT_MS, - OutputCaptureBuffer, REMOTE_BASH, REMOTE_WALK_TIMEOUT_MS, RefreshOutcome, optional_timeout, - resolve_path, validate_bash_probe, + OutputCaptureBuffer, OutputCaptureStats, REMOTE_BASH, REMOTE_WALK_TIMEOUT_MS, RefreshOutcome, + optional_timeout, resolve_path, validate_bash_probe, }; use crate::{ CommandOutputCallback, DirEntry, ExecResult, ExecStreamingRequest, ExecStreamingResult, @@ -2375,20 +2375,8 @@ impl Sandbox for DaytonaSandbox { .await?; } - let (stdout, stdout_capture) = { - let stdout_seen = stdout_seen.lock().await; - ( - String::from_utf8_lossy(&stdout_seen.to_bytes()).into_owned(), - stdout_seen.stats(), - ) - }; - let (stderr, stderr_capture) = { - let stderr_seen = stderr_seen.lock().await; - ( - String::from_utf8_lossy(&stderr_seen.to_bytes()).into_owned(), - stderr_seen.stats(), - ) - }; + let (stdout, stdout_capture) = drain_captured_stream(&stdout_seen).await; + let (stderr, stderr_capture) = drain_captured_stream(&stderr_seen).await; let result = ExecStreamingResult { result: ExecResult { @@ -2914,7 +2902,7 @@ async fn append_missing_log_suffix( } let mut seen = seen.lock().await; - let offset = captured_log_suffix_offset(&seen, final_bytes); + let offset = captured_log_suffix_offset(&mut seen, final_bytes); if offset >= final_bytes.len() { return Ok(()); } @@ -2928,21 +2916,34 @@ async fn append_missing_log_suffix( } } -fn captured_log_suffix_offset(seen: &OutputCaptureBuffer, final_bytes: &[u8]) -> usize { +/// Take the captured stream bytes out of their shared buffer as a lossy +/// string, avoiding a copy when the bytes are valid UTF-8. +async fn drain_captured_stream( + seen: &Arc>, +) -> (String, OutputCaptureStats) { + let buffer = { + let mut seen = seen.lock().await; + std::mem::replace(&mut *seen, OutputCaptureBuffer::new(None)) + }; + let (bytes, stats) = buffer.into_parts(); + let text = match String::from_utf8(bytes) { + Ok(text) => text, + Err(err) => String::from_utf8_lossy(err.as_bytes()).into_owned(), + }; + (text, stats) +} + +fn captured_log_suffix_offset(seen: &mut OutputCaptureBuffer, final_bytes: &[u8]) -> usize { let stats = seen.stats(); if stats.omitted_bytes == 0 { return missing_log_suffix_offset(&seen.to_bytes(), final_bytes); } - let observed_bytes = seen.observed_bytes(); - let head = seen.retained_head(); - let tail = seen.retained_tail(); + let observed_bytes = stats.observed_bytes; + let (head, tail) = seen.retained_slices(); if final_bytes.len() >= observed_bytes && final_bytes.starts_with(head) - && tail.iter().copied().eq(final_bytes - [observed_bytes.saturating_sub(tail.len())..observed_bytes] - .iter() - .copied()) + && tail == &final_bytes[observed_bytes.saturating_sub(tail.len())..observed_bytes] { return observed_bytes; } @@ -2952,12 +2953,7 @@ fn captured_log_suffix_offset(seen: &OutputCaptureBuffer, final_bytes: &[u8]) -> let max_overlap = tail.len().min(final_bytes.len()); for overlap in (1..=max_overlap).rev() { - if tail - .iter() - .skip(tail.len() - overlap) - .copied() - .eq(final_bytes[..overlap].iter().copied()) - { + if tail[tail.len() - overlap..] == final_bytes[..overlap] { return overlap; } } @@ -4791,9 +4787,9 @@ mod tests { let mut seen = OutputCaptureBuffer::new(Some(6)); seen.push(b"abcdefgh"); - assert_eq!(captured_log_suffix_offset(&seen, b"abcdefghij"), 8); - assert_eq!(captured_log_suffix_offset(&seen, b"abcdefgh"), 8); - assert_eq!(captured_log_suffix_offset(&seen, b"abcd"), 4); + assert_eq!(captured_log_suffix_offset(&mut seen, b"abcdefghij"), 8); + assert_eq!(captured_log_suffix_offset(&mut seen, b"abcdefgh"), 8); + assert_eq!(captured_log_suffix_offset(&mut seen, b"abcd"), 4); } #[test] diff --git a/lib/components/fabro-sandbox/src/sandbox.rs b/lib/components/fabro-sandbox/src/sandbox.rs index 473cfa71b..a08094df6 100644 --- a/lib/components/fabro-sandbox/src/sandbox.rs +++ b/lib/components/fabro-sandbox/src/sandbox.rs @@ -870,38 +870,34 @@ impl OutputCaptureBuffer { #[cfg(feature = "daytona")] #[must_use] pub(crate) fn to_bytes(&self) -> Vec { - let stats = self.stats(); - let mut bytes = Vec::with_capacity(stats.retained_bytes); + let mut bytes = Vec::with_capacity(self.head.len().saturating_add(self.tail.len())); bytes.extend_from_slice(&self.head); - bytes.extend(self.tail.iter().copied()); + let (front, back) = self.tail.as_slices(); + bytes.extend_from_slice(front); + bytes.extend_from_slice(back); bytes } #[must_use] pub(crate) fn into_parts(self) -> (Vec, OutputCaptureStats) { let stats = self.stats(); - let mut bytes = Vec::with_capacity(stats.retained_bytes); - bytes.extend(self.head); - bytes.extend(self.tail); + let Self { + head: mut bytes, + tail, + .. + } = self; + let (front, back) = tail.as_slices(); + bytes.extend_from_slice(front); + bytes.extend_from_slice(back); (bytes, stats) } + /// Retained bytes as two contiguous slices: the stable head, then the + /// rolling tail. #[cfg(feature = "daytona")] #[must_use] - pub(crate) fn observed_bytes(&self) -> usize { - self.observed_bytes - } - - #[cfg(feature = "daytona")] - #[must_use] - pub(crate) fn retained_head(&self) -> &[u8] { - &self.head - } - - #[cfg(feature = "daytona")] - #[must_use] - pub(crate) fn retained_tail(&self) -> &VecDeque { - &self.tail + pub(crate) fn retained_slices(&mut self) -> (&[u8], &[u8]) { + (&self.head, self.tail.make_contiguous()) } } @@ -1008,14 +1004,8 @@ pub(crate) async fn replay_exec_result( .await?; } } - let mut stdout_capture = OutputCaptureBuffer::new(stream_output_bytes_cap); - stdout_capture.push(result.stdout.as_bytes()); - let mut stderr_capture = OutputCaptureBuffer::new(stream_output_bytes_cap); - stderr_capture.push(result.stderr.as_bytes()); - let (stdout, stdout_capture) = stdout_capture.into_parts(); - let (stderr, stderr_capture) = stderr_capture.into_parts(); - result.stdout = String::from_utf8_lossy(&stdout).into_owned(); - result.stderr = String::from_utf8_lossy(&stderr).into_owned(); + let stdout_capture = capture_replayed_stream(&mut result.stdout, stream_output_bytes_cap); + let stderr_capture = capture_replayed_stream(&mut result.stderr, stream_output_bytes_cap); Ok(ExecStreamingResult { result, @@ -1026,6 +1016,21 @@ pub(crate) async fn replay_exec_result( }) } +/// Bound one replayed stream in place, leaving it untouched when it already +/// fits the cap. +fn capture_replayed_stream(text: &mut String, cap: Option) -> OutputCaptureStats { + match cap { + Some(cap) if text.len() > cap => { + let mut buffer = OutputCaptureBuffer::new(Some(cap)); + buffer.push(text.as_bytes()); + let (bytes, stats) = buffer.into_parts(); + *text = String::from_utf8_lossy(&bytes).into_owned(); + stats + } + _ => OutputCaptureStats::complete(text.len()), + } +} + pub struct StdioProcess { pub stdin: Pin>, pub stdout: Pin>, diff --git a/lib/foundation/fabro-types/src/run_event/mod.rs b/lib/foundation/fabro-types/src/run_event/mod.rs index 30d56a110..4dc457455 100644 --- a/lib/foundation/fabro-types/src/run_event/mod.rs +++ b/lib/foundation/fabro-types/src/run_event/mod.rs @@ -22,6 +22,14 @@ pub use todo::*; use crate::{ParallelBranchId, Principal, RunId, StageId}; +/// Maximum accepted body size for `POST /runs/{id}/events`. +/// +/// Producers that embed large payloads in an event (serialized tool output in +/// particular) must budget against this limit, leaving headroom for the rest +/// of the event envelope. The agent layer reserves half of it for serialized +/// tool output. +pub const MAX_RUN_EVENT_BODY_BYTES: usize = 3 * 1024 * 1024; + #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum RunNoticeLevel {