mirror of
https://github.com/aaif-goose/goose.git
synced 2026-07-03 14:10:03 +02:00
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:
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user