From 6e19fb2eec94c6799e34890af7a949245edbc862 Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Mon, 24 Aug 2026 13:00:58 -0400 Subject: [PATCH] Reserve event space for serialized tool output --- .../fabro-server/src/server/handler/events.rs | 7 +- lib/apps/fabro-server/src/server/tests.rs | 73 ++++++++++++++ .../fabro-agent/src/tool_execution.rs | 94 +++++++++++++++++++ lib/components/fabro-agent/src/truncation.rs | 69 ++++++++++++-- 4 files changed, 233 insertions(+), 10 deletions(-) diff --git a/lib/apps/fabro-server/src/server/handler/events.rs b/lib/apps/fabro-server/src/server/handler/events.rs index 926fefdca..247823cb7 100644 --- a/lib/apps/fabro-server/src/server/handler/events.rs +++ b/lib/apps/fabro-server/src/server/handler/events.rs @@ -1,5 +1,6 @@ use std::sync::Arc; +use axum::extract::DefaultBodyLimit; use fabro_types::{ RunEventDetailContent, RunEventDetailContentKind, RunEventDetailEnvelope, RunEventDetailResponse, @@ -15,12 +16,16 @@ 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)) .route( "/runs/{id}/events", - get(list_run_events).post(append_run_event), + get(list_run_events) + .post(append_run_event) + .layer(DefaultBodyLimit::max(MAX_RUN_EVENT_BODY_BYTES)), ) .route("/runs/{id}/events/{seq}", get(get_run_event_detail)) .route( diff --git a/lib/apps/fabro-server/src/server/tests.rs b/lib/apps/fabro-server/src/server/tests.rs index dc21a12e6..61429e6e1 100644 --- a/lib/apps/fabro-server/src/server/tests.rs +++ b/lib/apps/fabro-server/src/server/tests.rs @@ -10985,6 +10985,79 @@ async fn append_run_event_rejects_run_id_mismatch() { assert_status!(response, StatusCode::BAD_REQUEST).await; } +#[tokio::test] +async fn append_run_event_accepts_a_body_larger_than_two_mib() { + let state = test_app_state(); + let app = crate::test_support::build_test_router(Arc::clone(&state)); + let run_id = create_run(&app, MINIMAL_DOT).await; + let payload = json!({ + "id": "evt-large-agent-output", + "ts": "2026-08-24T12:00:00Z", + "run_id": run_id, + "event": "agent.tool.completed", + "properties": { + "tool_name": "shell", + "tool_call_id": "call-large", + "output": "x".repeat(2 * 1024 * 1024), + "is_error": false, + "visit": 1 + } + }) + .to_string(); + assert!(payload.len() > 2 * 1024 * 1024); + assert!(payload.len() < 3 * 1024 * 1024); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri(api(&format!("/runs/{run_id}/events"))) + .header("content-type", "application/json") + .body(Body::from(payload)) + .unwrap(), + ) + .await + .unwrap(); + + assert_status!(response, StatusCode::OK).await; +} + +#[tokio::test] +async fn append_run_event_rejects_a_body_larger_than_three_mib() { + let state = test_app_state(); + let app = crate::test_support::build_test_router(Arc::clone(&state)); + let run_id = create_run(&app, MINIMAL_DOT).await; + let payload = json!({ + "id": "evt-oversized-agent-output", + "ts": "2026-08-24T12:00:00Z", + "run_id": run_id, + "event": "agent.tool.completed", + "properties": { + "tool_name": "shell", + "tool_call_id": "call-oversized", + "output": "x".repeat(3 * 1024 * 1024), + "is_error": false, + "visit": 1 + } + }) + .to_string(); + assert!(payload.len() > 3 * 1024 * 1024); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri(api(&format!("/runs/{run_id}/events"))) + .header("content-type", "application/json") + .body(Body::from(payload)) + .unwrap(), + ) + .await + .unwrap(); + + assert_status!(response, StatusCode::PAYLOAD_TOO_LARGE).await; +} + #[tokio::test] async fn append_run_event_rejects_reserved_archive_event() { let state = test_app_state(); diff --git a/lib/components/fabro-agent/src/tool_execution.rs b/lib/components/fabro-agent/src/tool_execution.rs index 3cb8b674a..7b377e586 100644 --- a/lib/components/fabro-agent/src/tool_execution.rs +++ b/lib/components/fabro-agent/src/tool_execution.rs @@ -645,6 +645,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 tokio::sync::broadcast; use super::*; @@ -660,6 +661,7 @@ mod tests { use crate::test_support::MockSandbox; use crate::tool_registry::{RegisteredTool, ToolContext, ToolRegistry, ToolSource}; use crate::tools::make_shell_tool; + use crate::truncation::MAX_SERIALIZED_TOOL_OUTPUT_BYTES; use crate::types::SessionEvent; struct NamedPolicy { @@ -1026,6 +1028,98 @@ mod tests { assert!(event_output.contains(&format!("... {} bytes omitted ...", completed.4))); } + #[tokio::test] + async fn serialized_tool_output_and_full_event_stay_within_reserved_budgets() { + let mut registry = ToolRegistry::new(); + registry.register(make_echo_tool()); + let text = format!( + "HEAD{}TAIL", + "\0".repeat(MAX_RETAINED_TOOL_OUTPUT_BYTES - "echo: HEADTAIL".len()) + ); + let tc = make_tool_call("echo", "call_escaped", serde_json::json!({"text": text})); + let emitter = Emitter::new(); + let mut receiver = emitter.subscribe(); + + let result = execute_and_emit_one_tool( + &tc, + ®istry, + make_sandbox(), + None, + CancellationToken::new(), + &SessionOptions::default(), + &emitter, + "test-session", + "test-session", + None, + ) + .await; + + assert!(!result.is_error); + let completed = loop { + let event = receiver.try_recv().expect("tool completion event"); + if let AgentEvent::ToolCallCompleted { + tool_name, + tool_call_id, + output, + is_error, + output_bytes_observed, + output_bytes_retained, + output_bytes_omitted, + } = event.event + { + break ( + tool_name, + tool_call_id, + output, + is_error, + output_bytes_observed, + output_bytes_retained, + output_bytes_omitted, + ); + } + }; + + let serialized_output_bytes = serde_json::to_vec(&completed.2) + .expect("tool output serializes") + .len(); + assert!(serialized_output_bytes <= MAX_SERIALIZED_TOOL_OUTPUT_BYTES); + + let run_id = fabro_types::RunId::new(); + let run_event = fabro_types::RunEvent { + id: "evt-escaped-output".to_string(), + ts: chrono::Utc::now(), + run_id, + node_id: None, + node_label: None, + stage_id: None, + parallel_group_id: None, + parallel_branch_id: None, + session_id: Some("test-session".to_string()), + parent_session_id: None, + tool_call_id: Some(completed.1.clone()), + actor: None, + body: fabro_types::EventBody::AgentToolCompleted(AgentToolCompletedProps { + tool_name: completed.0, + tool_call_id: completed.1, + output: completed.2, + is_error: completed.3, + visit: 1, + output_bytes_observed: Some(completed.4 as u64), + output_bytes_retained: Some(completed.5 as u64), + output_bytes_omitted: Some(completed.6 as u64), + tool_result: None, + turn_id: None, + }), + }; + let serialized_event_bytes = serde_json::to_vec(&run_event) + .expect("run event serializes") + .len(); + assert!( + serialized_event_bytes < 2 * 1024 * 1024, + "serialized event was {serialized_event_bytes} bytes" + ); + } + #[tokio::test] async fn post_tool_use_hook_fires_on_success() { let mut registry = ToolRegistry::new(); diff --git a/lib/components/fabro-agent/src/truncation.rs b/lib/components/fabro-agent/src/truncation.rs index af85eb2c6..6e155da86 100644 --- a/lib/components/fabro-agent/src/truncation.rs +++ b/lib/components/fabro-agent/src/truncation.rs @@ -3,6 +3,7 @@ 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; #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct RetainedToolOutput { @@ -59,7 +60,7 @@ pub(crate) fn retain_tool_output( } /// Build the final model-facing preview, including truncation notices inside -/// the total byte budget. +/// the total byte budget and JSON serialization limit. #[must_use] pub(crate) fn preview_tool_output( output: &str, @@ -69,12 +70,16 @@ pub(crate) fn preview_tool_output( let mut content_budget = max_bytes; loop { let retained = retain_tool_output(output, content_budget, previously_omitted_bytes); - if retained.stats.omitted_bytes == 0 { - return retained; - } - - let rendered = render_retained_output(&retained); - if rendered.len() <= max_bytes { + let rendered = if retained.stats.omitted_bytes == 0 { + retained.output.clone() + } else { + render_retained_output(&retained) + }; + let serialized_bytes = serialized_json_string_bytes(&rendered); + 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, @@ -82,8 +87,20 @@ pub(crate) fn preview_tool_output( }; } - let excess = rendered.len().saturating_sub(max_bytes).max(1); - let next_budget = content_budget.saturating_sub(excess); + let mut next_budget = content_budget; + if rendered.len() > max_bytes { + let excess = rendered.len().saturating_sub(max_bytes); + next_budget = next_budget.min(content_budget.saturating_sub(excess)); + } + if serialized_bytes > MAX_SERIALIZED_TOOL_OUTPUT_BYTES { + let scaled_budget = (content_budget as u128) + .saturating_mul(MAX_SERIALIZED_TOOL_OUTPUT_BYTES as u128) + .checked_div(serialized_bytes as u128) + .and_then(|budget| usize::try_from(budget).ok()) + .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), @@ -95,6 +112,12 @@ pub(crate) fn preview_tool_output( } } +fn serialized_json_string_bytes(output: &str) -> usize { + serde_json::to_vec(output) + .expect("strings always serialize as JSON") + .len() +} + fn render_retained_output(retained: &RetainedToolOutput) -> String { let head = &retained.output[..retained.head_bytes]; let tail = &retained.output[retained.head_bytes..]; @@ -350,6 +373,34 @@ mod tests { assert!(preview.output.ends_with("efgh")); } + #[test] + fn model_preview_bounds_pathological_json_serialization() { + let output = format!( + "HEAD{}TAIL", + "\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); + + let preview = preview_tool_output(&output, MAX_RETAINED_TOOL_OUTPUT_BYTES, 0); + let serialized_bytes = serialized_json_string_bytes(&preview.output); + + assert!(preview.output.len() <= MAX_RETAINED_TOOL_OUTPUT_BYTES); + assert!( + serialized_bytes <= MAX_SERIALIZED_TOOL_OUTPUT_BYTES, + "serialized preview was {serialized_bytes} bytes" + ); + assert!(preview.output.starts_with("Warning: truncated output")); + assert!(preview.output.contains("HEAD")); + assert!(preview.output.ends_with("TAIL")); + assert_eq!(preview.stats.observed_bytes, MAX_RETAINED_TOOL_OUTPUT_BYTES); + assert!(preview.stats.retained_bytes < MAX_RETAINED_TOOL_OUTPUT_BYTES); + assert_eq!( + preview.stats.omitted_bytes, + preview.stats.observed_bytes - preview.stats.retained_bytes + ); + } + #[test] fn under_limit_passthrough_chars() { let output = "short output";