BRAINROT P2: per-field cap inside code_execution result bundles

The code_execution sandbox encourages the model to bundle multiple
sub-tool calls and return one labeled object like
'Result: { foo, bar, baz }'. Without per-field truncation, a single fat
field (full source file, build log, large query result) bloats the
whole bundle even when the other fields are tiny and useful.

Add truncate_execute_output() which intercepts pctx_code_mode's
ExecuteOutput just before output.markdown() is called and:

  - caps stdout and stderr if they exceed the per-field cap,
  - walks the  JSON value tree and truncates any leaf string
    longer than the cap with a marker.

Truncation is UTF-8 safe and recursive (nested objects + arrays). The
top-level P0 cap (16 KB on the whole tool_result text) still applies as
a backstop; this P2 cap (8 KB per field by default) operates earlier so
that one large field can't drown out the small useful ones in the same
bundle.

- Defaults: 8192 bytes per string field.
- Tunable: GOOSE_CODE_EXEC_FIELD_MAX_BYTES (0 to disable).
- Applied to both handle_execute_typescript and handle_execute_bash
  (execute_bash also wraps a bash invocation whose stdout/stderr can be
  large; same treatment).

Three serial unit tests cover (a) leaf-only truncation across object,
array, and nested object, (b) end-to-end ExecuteOutput cap, (c) env
disable.

Signed-off-by: Michael Neale <michael.neale@gmail.com>
This commit is contained in:
Michael Neale
2026-05-20 19:10:57 +10:00
parent a3d0d02c79
commit fb1119cc02
@@ -28,6 +28,89 @@ use tokio_util::sync::CancellationToken;
pub static EXTENSION_NAME: &str = "code_execution";
/// Environment variable for capping individual string fields inside the
/// bundled result of `execute_typescript` / `execute_bash`. Value is in
/// bytes per field. Set to `0` to disable. Defaults to
/// [`DEFAULT_CODE_EXEC_FIELD_MAX_BYTES`].
pub const CODE_EXEC_FIELD_MAX_BYTES_ENV: &str = "GOOSE_CODE_EXEC_FIELD_MAX_BYTES";
/// Default per-string-field cap inside a code-execution result bundle. Picked
/// smaller than [`crate::conversation::message::DEFAULT_TOOL_RESULT_MAX_BYTES`]
/// because bundles can hold many fields — one fat field shouldn't crowd out
/// the small useful ones in the same `{ ... }` object.
pub const DEFAULT_CODE_EXEC_FIELD_MAX_BYTES: usize = 8_192;
fn code_exec_field_cap() -> Option<usize> {
match std::env::var(CODE_EXEC_FIELD_MAX_BYTES_ENV) {
Ok(v) => v
.trim()
.parse::<usize>()
.ok()
.and_then(|n| if n == 0 { None } else { Some(n) }),
Err(_) => Some(DEFAULT_CODE_EXEC_FIELD_MAX_BYTES),
}
}
/// Walk a JSON value and replace any leaf string longer than `cap` bytes with
/// a marker telling the model how many bytes were elided. UTF-8 safe.
fn truncate_json_strings(value: &mut serde_json::Value, cap: usize) {
match value {
serde_json::Value::String(s) => {
if s.len() > cap {
let mut end = cap;
while end > 0 && !s.is_char_boundary(end) {
end -= 1;
}
let elided = s.len() - end;
let truncated = s.get(..end).unwrap_or("");
let marker = format!(
"{truncated}\n[goose: truncated {elided} bytes from this field (per-field cap={cap} bytes, set {env}=0 to disable). Re-run with a narrower selection (e.g. line/limit, head, grep) or return only the slice you need.]",
env = CODE_EXEC_FIELD_MAX_BYTES_ENV,
);
*s = marker;
}
}
serde_json::Value::Array(arr) => {
for item in arr.iter_mut() {
truncate_json_strings(item, cap);
}
}
serde_json::Value::Object(map) => {
for (_, v) in map.iter_mut() {
truncate_json_strings(v, cap);
}
}
_ => {}
}
}
/// Apply per-field truncation to an `ExecuteOutput` in place, capping
/// `stdout`, `stderr`, and any string leaves inside `output`.
fn truncate_execute_output(out: &mut pctx_code_mode::model::ExecuteOutput) {
let cap = match code_exec_field_cap() {
Some(c) => c,
None => return,
};
if out.stdout.len() > cap {
let mut s = std::mem::take(&mut out.stdout);
let mut tmp = serde_json::Value::String(std::mem::take(&mut s));
truncate_json_strings(&mut tmp, cap);
if let serde_json::Value::String(new_s) = tmp {
out.stdout = new_s;
}
}
if out.stderr.len() > cap {
let mut tmp = serde_json::Value::String(std::mem::take(&mut out.stderr));
truncate_json_strings(&mut tmp, cap);
if let serde_json::Value::String(new_s) = tmp {
out.stderr = new_s;
}
}
if let Some(v) = out.output.as_mut() {
truncate_json_strings(v, cap);
}
}
/// Append goose-specific context-hygiene advice to the upstream
/// `pctx_code_mode` tool descriptions. The upstream descriptions push the
/// model toward bundling many sub-tool calls into one execute and returning a
@@ -249,6 +332,8 @@ impl CodeExecutionClient {
.await
.map_err(|e| format!("Typescript execution task failed: {e}"))??;
let mut output = output;
truncate_execute_output(&mut output);
Ok(vec![Content::text(output.markdown())])
}
@@ -287,6 +372,8 @@ impl CodeExecutionClient {
.await
.map_err(|e| format!("Typescript execution task failed: {e}"))??;
let mut output = output;
truncate_execute_output(&mut output);
Ok(vec![Content::text(output.markdown())])
}
}
@@ -577,6 +664,8 @@ impl CodeModeState {
#[cfg(test)]
mod tests {
use super::*;
use pctx_code_mode::model::ExecuteOutput;
use serial_test::serial;
#[test]
fn context_hygiene_advice_is_appended_to_upstream_desc() {
@@ -587,4 +676,73 @@ mod tests {
assert!(combined.contains("individually small"));
assert!(combined.contains("Return only the fields you actually need"));
}
#[test]
#[serial]
fn truncate_json_strings_caps_long_leaves_only() {
std::env::remove_var(CODE_EXEC_FIELD_MAX_BYTES_ENV);
let mut v = serde_json::json!({
"small": "hello",
"big": "x".repeat(20_000),
"nested": {
"deep_big": "y".repeat(10_000),
"deep_small": "ok",
},
"arr": ["a", "z".repeat(20_000)],
"n": 5,
});
truncate_json_strings(&mut v, 1024);
assert_eq!(v["small"], "hello");
assert!(v["big"].as_str().unwrap().contains("[goose: truncated"));
assert!(v["big"].as_str().unwrap().len() < 2_000);
assert!(v["nested"]["deep_big"]
.as_str()
.unwrap()
.contains("[goose: truncated"));
assert_eq!(v["nested"]["deep_small"], "ok");
assert_eq!(v["arr"][0], "a");
assert!(v["arr"][1].as_str().unwrap().contains("[goose: truncated"));
assert_eq!(v["n"], 5);
}
#[test]
#[serial]
fn truncate_execute_output_caps_stdout_stderr_and_output_fields() {
std::env::set_var(CODE_EXEC_FIELD_MAX_BYTES_ENV, "2048");
let mut out = ExecuteOutput {
success: true,
stdout: "o".repeat(10_000),
stderr: "e".repeat(10_000),
output: Some(serde_json::json!({
"foo": "f".repeat(10_000),
"bar": "small",
})),
};
truncate_execute_output(&mut out);
assert!(out.stdout.len() < 5_000);
assert!(out.stdout.contains("[goose: truncated"));
assert!(out.stderr.len() < 5_000);
assert!(out.stderr.contains("[goose: truncated"));
let obj = out.output.as_ref().unwrap().as_object().unwrap();
assert!(obj["foo"].as_str().unwrap().contains("[goose: truncated"));
assert_eq!(obj["bar"], "small");
std::env::remove_var(CODE_EXEC_FIELD_MAX_BYTES_ENV);
}
#[test]
#[serial]
fn truncate_execute_output_disabled_via_env_zero() {
std::env::set_var(CODE_EXEC_FIELD_MAX_BYTES_ENV, "0");
let big = "x".repeat(50_000);
let mut out = ExecuteOutput {
success: true,
stdout: big.clone(),
stderr: String::new(),
output: Some(serde_json::json!({ "foo": big.clone() })),
};
truncate_execute_output(&mut out);
assert_eq!(out.stdout, big);
assert_eq!(out.output.unwrap()["foo"], big);
std::env::remove_var(CODE_EXEC_FIELD_MAX_BYTES_ENV);
}
}