Skip to main content

lash_plugin_tool_output_budget/
lib.rs

1use std::fs;
2use std::path::{Path, PathBuf};
3use std::sync::Arc;
4
5use sha2::{Digest, Sha256};
6
7use lash_core::plugin::{
8    PluginError, PluginFactory, PluginRegistrar, PluginSessionContext, SessionPlugin,
9    ToolResultProjectionContext,
10};
11use lash_core::{ModelToolReturn, ModelToolReturnPart, PluginStack, ToolCallOutcome, ToolValue};
12
13const APPROX_BYTES_PER_TOKEN: usize = 4;
14pub const DEFAULT_TOOL_OUTPUT_BUDGET_LIMIT_BYTES: usize = 16 * 1024;
15pub const DEFAULT_TOOL_OUTPUT_BUDGET_MAX_LINES: usize = 400;
16
17#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
18#[serde(rename_all = "snake_case")]
19pub enum ToolOutputBudgetMode {
20    Bytes,
21    Tokens,
22}
23
24#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
25#[serde(default)]
26pub struct ToolOutputBudgetConfig {
27    pub mode: ToolOutputBudgetMode,
28    pub limit: usize,
29    pub max_lines: usize,
30}
31
32impl Default for ToolOutputBudgetConfig {
33    fn default() -> Self {
34        Self {
35            mode: ToolOutputBudgetMode::Bytes,
36            limit: DEFAULT_TOOL_OUTPUT_BUDGET_LIMIT_BYTES,
37            max_lines: DEFAULT_TOOL_OUTPUT_BUDGET_MAX_LINES,
38        }
39    }
40}
41
42#[derive(Clone, Copy, Debug, PartialEq, Eq)]
43enum ProjectionDirection {
44    Head,
45    Tail,
46}
47
48/// Which end of the output a windowed truncation keeps.
49///
50/// `Head` keeps the leading lines (the common case); `Tail` keeps the
51/// trailing lines (used for streaming command output where the end is
52/// the interesting part). This is the public mirror of the budget
53/// plugin's internal `ProjectionDirection`, exported so other plugins
54/// (e.g. rolling-history) can share the one canonical truncation core
55/// instead of forking it.
56#[derive(Clone, Copy, Debug, PartialEq, Eq)]
57pub enum TruncationDirection {
58    Head,
59    Tail,
60}
61
62impl From<ProjectionDirection> for TruncationDirection {
63    fn from(direction: ProjectionDirection) -> Self {
64        match direction {
65            ProjectionDirection::Head => TruncationDirection::Head,
66            ProjectionDirection::Tail => TruncationDirection::Tail,
67        }
68    }
69}
70
71/// The unit reported in the `...N <unit> truncated...` marker when a
72/// windowed truncation hits its byte budget (rather than its line cap).
73#[derive(Clone, Copy, Debug, PartialEq, Eq)]
74pub enum TruncationUnit {
75    /// Report removed *characters*, labelled `bytes`.
76    Bytes,
77    /// Report an approximate removed *token* count, labelled `tokens`.
78    Tokens,
79}
80
81impl TruncationUnit {
82    fn label(self) -> &'static str {
83        match self {
84            TruncationUnit::Bytes => "bytes",
85            TruncationUnit::Tokens => "tokens",
86        }
87    }
88}
89
90/// Parameters for [`truncate_windowed`], the single canonical
91/// head/tail-window + byte-cap truncation implementation shared by the
92/// tool-output-budget projector and the rolling-history compaction
93/// plugin.
94#[derive(Clone, Copy, Debug)]
95pub struct WindowedTruncation<'a> {
96    /// Maximum number of lines retained in the preview window.
97    pub max_lines: usize,
98    /// Maximum number of bytes retained in the preview window.
99    pub max_bytes: usize,
100    /// Which end of the output to keep.
101    pub direction: TruncationDirection,
102    /// The unit reported in the byte-budget truncation marker.
103    pub unit: TruncationUnit,
104    /// Trailing hint text appended to (Head) / prepended to (Tail) the
105    /// preview, explaining the truncation and where the full output is.
106    pub hint: &'a str,
107}
108
109/// The canonical head/tail-window + byte-cap truncation core.
110///
111/// Returns `text` unchanged when it already fits within `max_lines` and
112/// `max_bytes`. Otherwise keeps a preview window from the configured end
113/// and wraps it with a `...N <unit> truncated...` marker plus the
114/// caller-supplied `hint`.
115///
116/// A single line that is itself larger than `max_bytes` is truncated at
117/// a UTF-8 char boundary rather than dropped, so over-long lines never
118/// silently disappear and the function never panics on multi-byte text.
119pub fn truncate_windowed(text: &str, opts: &WindowedTruncation) -> String {
120    let lines: Vec<&str> = text.lines().collect();
121    let total_bytes = text.len();
122    if lines.len() <= opts.max_lines && total_bytes <= opts.max_bytes {
123        return text.to_string();
124    }
125
126    let mut preview_lines: Vec<String> = Vec::new();
127    let mut bytes = 0usize;
128    let mut hit_budget = false;
129
130    let mut push_line = |line: &str, bytes: &mut usize, hit_budget: &mut bool| -> bool {
131        // `separator` accounts for the `\n` re-joined between lines; the
132        // first retained line carries no separator.
133        let separator = usize::from(!preview_lines.is_empty());
134        let remaining = opts.max_bytes.saturating_sub(*bytes + separator);
135        if line.len() + separator <= opts.max_bytes.saturating_sub(*bytes) {
136            preview_lines.push(line.to_string());
137            *bytes += line.len() + separator;
138            true
139        } else if preview_lines.is_empty() && remaining > 0 {
140            // A lone line longer than the whole budget: truncate it at a
141            // char boundary instead of dropping it entirely.
142            let cut = char_floor(line, remaining);
143            if cut == 0 {
144                *hit_budget = true;
145                return false;
146            }
147            preview_lines.push(line[..cut].to_string());
148            *bytes += cut;
149            *hit_budget = true;
150            false
151        } else {
152            *hit_budget = true;
153            false
154        }
155    };
156
157    match opts.direction {
158        TruncationDirection::Head => {
159            for line in lines.iter().take(opts.max_lines) {
160                if !push_line(line, &mut bytes, &mut hit_budget) {
161                    break;
162                }
163            }
164        }
165        TruncationDirection::Tail => {
166            for line in lines.iter().rev().take(opts.max_lines) {
167                if !push_line(line, &mut bytes, &mut hit_budget) {
168                    break;
169                }
170            }
171            preview_lines.reverse();
172        }
173    }
174
175    let preview = preview_lines.join("\n");
176    let (removed, unit) = if hit_budget {
177        let removed = match opts.unit {
178            TruncationUnit::Bytes => {
179                u64::try_from(text.chars().count().saturating_sub(preview.chars().count()))
180                    .unwrap_or(u64::MAX)
181            }
182            TruncationUnit::Tokens => {
183                approx_tokens_from_byte_count(total_bytes.saturating_sub(preview.len()))
184            }
185        };
186        (removed, opts.unit.label())
187    } else {
188        (
189            u64::try_from(lines.len().saturating_sub(preview_lines.len())).unwrap_or(u64::MAX),
190            "lines",
191        )
192    };
193    let hint = opts.hint;
194    match opts.direction {
195        TruncationDirection::Head => {
196            format!("{preview}\n\n...{removed} {unit} truncated...\n\n{hint}")
197        }
198        TruncationDirection::Tail => {
199            format!("...{removed} {unit} truncated...\n\n{hint}\n\n{preview}")
200        }
201    }
202}
203
204/// Largest byte offset `<= max` that lands on a UTF-8 char boundary.
205fn char_floor(text: &str, max: usize) -> usize {
206    if max >= text.len() {
207        return text.len();
208    }
209    let mut cut = max;
210    while cut > 0 && !text.is_char_boundary(cut) {
211        cut -= 1;
212    }
213    cut
214}
215
216pub struct ToolOutputBudgetPluginFactory {
217    config: ToolOutputBudgetConfig,
218}
219
220impl ToolOutputBudgetPluginFactory {
221    pub fn new(config: ToolOutputBudgetConfig) -> Self {
222        Self { config }
223    }
224}
225
226impl Default for ToolOutputBudgetPluginFactory {
227    fn default() -> Self {
228        Self::new(ToolOutputBudgetConfig::default())
229    }
230}
231
232pub fn tool_output_budget_stack() -> PluginStack {
233    let mut stack = PluginStack::new();
234    stack.push(Arc::new(ToolOutputBudgetPluginFactory::default()));
235    stack
236}
237
238impl PluginFactory for ToolOutputBudgetPluginFactory {
239    fn id(&self) -> &'static str {
240        "tool_output_budget"
241    }
242
243    fn build(&self, _ctx: &PluginSessionContext) -> Result<Arc<dyn SessionPlugin>, PluginError> {
244        Ok(Arc::new(ToolOutputBudgetPlugin {
245            config: self.config.clone(),
246        }))
247    }
248}
249
250struct ToolOutputBudgetPlugin {
251    config: ToolOutputBudgetConfig,
252}
253
254impl SessionPlugin for ToolOutputBudgetPlugin {
255    fn id(&self) -> &'static str {
256        "tool_output_budget"
257    }
258
259    fn register(&self, reg: &mut PluginRegistrar) -> Result<(), PluginError> {
260        register_projector(reg, &self.config)
261    }
262}
263
264fn register_projector(
265    reg: &mut PluginRegistrar,
266    config: &ToolOutputBudgetConfig,
267) -> Result<(), PluginError> {
268    let config = config.clone();
269    reg.tool_results().projector(Arc::new(move |ctx| {
270        let config = config.clone();
271        Box::pin(async move { Ok(project_tool_result(&config, ctx)) })
272    }))
273}
274
275fn project_tool_result(
276    config: &ToolOutputBudgetConfig,
277    ctx: ToolResultProjectionContext,
278) -> ModelToolReturn {
279    let parts = project_model_parts(config, &ctx);
280    ModelToolReturn {
281        call_id: ctx.call_id.clone(),
282        tool_name: ctx.tool_name.clone(),
283        parts,
284    }
285}
286
287/// Project a tool result into the rendered model-facing text using the
288/// canonical projector — including the structure-aware `batch` path that
289/// recurses into each child result with the correct per-child truncation
290/// direction.
291///
292/// Exposed so other plugins (e.g. rolling-history) can reuse the one
293/// canonical batch/structured projection instead of flattening batched
294/// output to an opaque string and tail-truncating it (which cuts JSON
295/// mid-structure and loses per-child windowing). Attachments are rendered
296/// inline using the same `[Attachment: …]` placeholder the budget plugin
297/// uses internally.
298pub fn project_tool_result_text(
299    config: &ToolOutputBudgetConfig,
300    ctx: ToolResultProjectionContext,
301) -> String {
302    render_model_return_parts(&project_tool_result(config, ctx).parts)
303}
304
305fn project_model_parts(
306    config: &ToolOutputBudgetConfig,
307    ctx: &ToolResultProjectionContext,
308) -> Vec<ModelToolReturnPart> {
309    if ctx.tool_name == "batch" {
310        let value = project_batch_value(config, ctx);
311        return vec![ModelToolReturnPart::text(render_projected_model_value(
312            &value,
313        ))];
314    }
315
316    match &ctx.output.outcome {
317        ToolCallOutcome::Success(value) => project_tool_value_parts(config, ctx, value),
318        ToolCallOutcome::Failure(failure) => {
319            let mut parts = vec![ModelToolReturnPart::text(
320                lash_core::session_model::format_tool_output_content(&ctx.output),
321            )];
322            if let Some(raw) = &failure.raw {
323                parts.extend(
324                    raw.attachments()
325                        .into_iter()
326                        .map(ModelToolReturnPart::Attachment),
327                );
328            }
329            parts
330        }
331        ToolCallOutcome::Cancelled(cancellation) => {
332            let mut parts = vec![ModelToolReturnPart::text(
333                lash_core::session_model::format_tool_output_content(&ctx.output),
334            )];
335            if let Some(raw) = &cancellation.raw {
336                parts.extend(
337                    raw.attachments()
338                        .into_iter()
339                        .map(ModelToolReturnPart::Attachment),
340                );
341            }
342            parts
343        }
344    }
345}
346
347fn render_projected_model_value(value: &serde_json::Value) -> String {
348    match value {
349        serde_json::Value::String(text) => text.clone(),
350        other => serde_json::to_string(other).unwrap_or_else(|_| "null".to_string()),
351    }
352}
353
354fn project_tool_value_parts(
355    config: &ToolOutputBudgetConfig,
356    ctx: &ToolResultProjectionContext,
357    value: &ToolValue,
358) -> Vec<ModelToolReturnPart> {
359    let mut parts = Vec::new();
360    match value {
361        ToolValue::String(text) => {
362            parts.push(ModelToolReturnPart::text(project_text(text, config, ctx)))
363        }
364        ToolValue::Attachment(reference) => {
365            parts.push(ModelToolReturnPart::Attachment(reference.clone()));
366        }
367        ToolValue::Null
368        | ToolValue::Bool(_)
369        | ToolValue::Number(_)
370        | ToolValue::Array(_)
371        | ToolValue::Object(_) => {
372            push_projected_tool_value_parts(value, &mut parts, config, ctx);
373        }
374    }
375    parts
376}
377
378fn push_projected_tool_value_parts(
379    value: &ToolValue,
380    parts: &mut Vec<ModelToolReturnPart>,
381    config: &ToolOutputBudgetConfig,
382    ctx: &ToolResultProjectionContext,
383) {
384    match value {
385        ToolValue::Null => push_text_part(parts, "null"),
386        ToolValue::Bool(value) => push_text_part(parts, value.to_string()),
387        ToolValue::Number(value) => push_text_part(parts, value.to_string()),
388        ToolValue::String(text) => push_text_part(
389            parts,
390            serde_json::to_string(&project_text(text, config, ctx))
391                .unwrap_or_else(|_| "\"\"".to_string()),
392        ),
393        ToolValue::Attachment(reference) => {
394            parts.push(ModelToolReturnPart::Attachment(reference.clone()));
395        }
396        ToolValue::Array(items) => {
397            push_text_part(parts, "[");
398            for (index, item) in items.iter().enumerate() {
399                if index > 0 {
400                    push_text_part(parts, ",");
401                }
402                push_projected_tool_value_parts(item, parts, config, ctx);
403            }
404            push_text_part(parts, "]");
405        }
406        ToolValue::Object(map) => {
407            push_text_part(parts, "{");
408            for (index, (key, value)) in map.iter().enumerate() {
409                if index > 0 {
410                    push_text_part(parts, ",");
411                }
412                push_text_part(
413                    parts,
414                    serde_json::to_string(key).unwrap_or_else(|_| "\"\"".to_string()),
415                );
416                push_text_part(parts, ":");
417                push_projected_tool_value_parts(value, parts, config, ctx);
418            }
419            push_text_part(parts, "}");
420        }
421    }
422}
423
424fn push_text_part(parts: &mut Vec<ModelToolReturnPart>, text: impl Into<String>) {
425    let text = text.into();
426    if text.is_empty() {
427        return;
428    }
429    if let Some(ModelToolReturnPart::Text { text: existing }) = parts.last_mut() {
430        existing.push_str(&text);
431    } else {
432        parts.push(ModelToolReturnPart::text(text));
433    }
434}
435
436fn project_text(
437    text: &str,
438    config: &ToolOutputBudgetConfig,
439    ctx: &ToolResultProjectionContext,
440) -> String {
441    if !needs_truncation(text, config) {
442        return text.to_string();
443    }
444    truncate_text(
445        text,
446        config,
447        tool_projection_direction(&ctx.tool_name),
448        Some(ctx),
449    )
450}
451
452fn needs_truncation(text: &str, config: &ToolOutputBudgetConfig) -> bool {
453    if text.lines().count() > config.max_lines {
454        return true;
455    }
456    match config.mode {
457        ToolOutputBudgetMode::Bytes => text.len() > config.limit,
458        ToolOutputBudgetMode::Tokens => approx_token_count(text) > config.limit,
459    }
460}
461
462fn truncate_text(
463    text: &str,
464    config: &ToolOutputBudgetConfig,
465    direction: ProjectionDirection,
466    ctx: Option<&ToolResultProjectionContext>,
467) -> String {
468    truncate_text_with_hint(text, config, direction, truncation_hint(ctx, text))
469}
470
471fn truncate_text_with_hint(
472    text: &str,
473    config: &ToolOutputBudgetConfig,
474    direction: ProjectionDirection,
475    hint: String,
476) -> String {
477    if text.is_empty() {
478        return String::new();
479    }
480    let max_bytes = match config.mode {
481        ToolOutputBudgetMode::Bytes => config.limit,
482        ToolOutputBudgetMode::Tokens => approx_bytes_for_tokens(config.limit),
483    };
484    if max_bytes == 0 {
485        return format_truncation_marker(
486            config.mode,
487            removed_units(config.mode, text.len(), text.chars().count()),
488        );
489    }
490    if !needs_truncation(text, config) {
491        return text.to_string();
492    }
493    truncate_windowed(
494        text,
495        &WindowedTruncation {
496            max_lines: config.max_lines,
497            max_bytes,
498            direction: direction.into(),
499            unit: match config.mode {
500                ToolOutputBudgetMode::Bytes => TruncationUnit::Bytes,
501                ToolOutputBudgetMode::Tokens => TruncationUnit::Tokens,
502            },
503            hint: &hint,
504        },
505    )
506}
507
508fn format_truncation_marker(mode: ToolOutputBudgetMode, removed: u64) -> String {
509    match mode {
510        ToolOutputBudgetMode::Bytes => format!("…{removed} chars truncated…"),
511        ToolOutputBudgetMode::Tokens => format!("…{removed} tokens truncated…"),
512    }
513}
514
515fn removed_units(mode: ToolOutputBudgetMode, removed_bytes: usize, removed_chars: usize) -> u64 {
516    match mode {
517        ToolOutputBudgetMode::Bytes => u64::try_from(removed_chars).unwrap_or(u64::MAX),
518        ToolOutputBudgetMode::Tokens => approx_tokens_from_byte_count(removed_bytes),
519    }
520}
521
522fn approx_token_count(text: &str) -> usize {
523    text.len()
524        .saturating_add(APPROX_BYTES_PER_TOKEN.saturating_sub(1))
525        / APPROX_BYTES_PER_TOKEN
526}
527
528fn approx_bytes_for_tokens(tokens: usize) -> usize {
529    tokens.saturating_mul(APPROX_BYTES_PER_TOKEN)
530}
531
532fn approx_tokens_from_byte_count(bytes: usize) -> u64 {
533    let bytes = bytes as u64;
534    bytes.saturating_add((APPROX_BYTES_PER_TOKEN as u64).saturating_sub(1))
535        / (APPROX_BYTES_PER_TOKEN as u64)
536}
537
538fn tool_projection_direction(tool_name: &str) -> ProjectionDirection {
539    match tool_name {
540        "exec_command" | "write_stdin" => ProjectionDirection::Tail,
541        _ => ProjectionDirection::Head,
542    }
543}
544
545fn truncation_hint(ctx: Option<&ToolResultProjectionContext>, text: &str) -> String {
546    let output_path = ctx
547        .and_then(existing_tool_output_path)
548        .or_else(|| ctx.and_then(|ctx| spill_tool_output(&ctx.tool_name, &ctx.args, text)));
549    match output_path {
550        Some(path) => format!(
551            "The tool output was truncated. Full output saved to: {}\nUse `read_file` with `offset`/`limit` or `grep` to inspect specific sections instead of reading the whole file at once.",
552            path.display()
553        ),
554        None => "The tool output was truncated. Use `read_file` with `offset`/`limit` or `grep` to inspect specific sections instead of reading the whole file at once.".to_string(),
555    }
556}
557
558fn existing_tool_output_path(ctx: &ToolResultProjectionContext) -> Option<PathBuf> {
559    ctx.output
560        .value_for_projection()
561        .get("full_output_path")
562        .and_then(|value| value.as_str())
563        .filter(|value| !value.trim().is_empty())
564        .map(PathBuf::from)
565}
566
567fn spill_tool_output(
568    tool_name: &str,
569    args: &serde_json::Value,
570    full_output: &str,
571) -> Option<PathBuf> {
572    let dir = std::env::temp_dir().join("lash-tool-output");
573    if fs::create_dir_all(&dir).is_err() {
574        return None;
575    }
576
577    let mut hasher = Sha256::new();
578    hasher.update(tool_name.as_bytes());
579    hasher.update(args.to_string().as_bytes());
580    hasher.update(full_output.as_bytes());
581    let digest = format!("{:x}", hasher.finalize());
582    let stem = tool_name
583        .chars()
584        .map(|ch| {
585            if ch.is_ascii_alphanumeric() || matches!(ch, '.' | '_' | '-') {
586                ch
587            } else {
588                '_'
589            }
590        })
591        .collect::<String>();
592    let path = dir.join(format!("{stem}-{}.txt", &digest[..12]));
593    if write_if_changed(&path, full_output).is_err() {
594        return None;
595    }
596    Some(path)
597}
598
599fn write_if_changed(path: &Path, content: &str) -> std::io::Result<()> {
600    let should_write = match fs::read_to_string(path) {
601        Ok(existing) => existing != content,
602        Err(_) => true,
603    };
604    if should_write {
605        fs::write(path, content)?;
606    }
607    Ok(())
608}
609
610fn project_batch_value(
611    config: &ToolOutputBudgetConfig,
612    ctx: &ToolResultProjectionContext,
613) -> serde_json::Value {
614    let value = ctx.output.value_for_projection();
615    let Some(map) = value.as_object() else {
616        return project_json_value(&value, config, ctx);
617    };
618
619    let mut projected = serde_json::Map::new();
620
621    let results = map
622        .get("results")
623        .and_then(|value| value.as_array())
624        .map(|items| {
625            items
626                .iter()
627                .enumerate()
628                .map(|(index, item)| project_batch_child_value(index, item, config, ctx))
629                .collect::<Vec<_>>()
630        })
631        .unwrap_or_default();
632    projected.insert("results".to_string(), serde_json::Value::Array(results));
633    serde_json::Value::Object(projected)
634}
635
636fn project_batch_child_value(
637    index: usize,
638    item: &serde_json::Value,
639    config: &ToolOutputBudgetConfig,
640    ctx: &ToolResultProjectionContext,
641) -> serde_json::Value {
642    let Some(map) = item.as_object() else {
643        return project_json_value(item, config, ctx);
644    };
645
646    let tool_name = map
647        .get("tool")
648        .and_then(|value| value.as_str())
649        .or_else(|| batch_child_tool_name(&ctx.args, index))
650        .unwrap_or("tool")
651        .to_string();
652    let success = map
653        .get("success")
654        .and_then(|value| value.as_bool())
655        .unwrap_or(false);
656    let duration_ms = map
657        .get("duration_ms")
658        .and_then(|value| value.as_u64())
659        .unwrap_or_default();
660    let child_value = if success {
661        map.get("result")
662            .cloned()
663            .unwrap_or(serde_json::Value::Null)
664    } else {
665        map.get("error").cloned().unwrap_or(serde_json::Value::Null)
666    };
667    let child_args = batch_child_args(&ctx.args, index);
668
669    let projected_child = if tool_name == "batch" || !success {
670        project_json_value(&child_value, config, ctx)
671    } else {
672        let model_return = project_tool_result(
673            config,
674            ToolResultProjectionContext {
675                session_id: ctx.session_id.clone(),
676                call_id: format!("{}.{}", ctx.call_id, index),
677                tool_name: tool_name.clone(),
678                args: child_args,
679                output: lash_core::ToolCallOutput::success(child_value.clone()),
680                duration_ms,
681            },
682        );
683        let rendered = render_model_return_parts(&model_return.parts);
684        rendered
685            .parse::<serde_json::Value>()
686            .unwrap_or(serde_json::Value::String(rendered))
687    };
688
689    let mut projected = serde_json::Map::new();
690    if let Some(value) = map.get("index") {
691        projected.insert("index".to_string(), value.clone());
692    }
693    projected.insert("tool".to_string(), serde_json::json!(tool_name));
694    projected.insert("success".to_string(), serde_json::json!(success));
695    projected.insert("duration_ms".to_string(), serde_json::json!(duration_ms));
696    projected.insert(
697        if success {
698            "result".to_string()
699        } else {
700            "error".to_string()
701        },
702        projected_child,
703    );
704    serde_json::Value::Object(projected)
705}
706
707fn render_model_return_parts(parts: &[ModelToolReturnPart]) -> String {
708    let mut rendered = String::new();
709    for part in parts {
710        match part {
711            ModelToolReturnPart::Text { text } => rendered.push_str(text),
712            ModelToolReturnPart::Attachment(reference) => {
713                rendered.push_str("[Attachment: ");
714                rendered.push_str(
715                    reference
716                        .label
717                        .as_deref()
718                        .unwrap_or_else(|| reference.id.as_str()),
719                );
720                rendered.push(']');
721            }
722        }
723    }
724    rendered
725}
726
727fn project_json_value(
728    value: &serde_json::Value,
729    config: &ToolOutputBudgetConfig,
730    ctx: &ToolResultProjectionContext,
731) -> serde_json::Value {
732    match value {
733        serde_json::Value::String(text) => {
734            serde_json::Value::String(project_text(text, config, ctx))
735        }
736        serde_json::Value::Array(items) => serde_json::Value::Array(
737            items
738                .iter()
739                .map(|item| project_json_value(item, config, ctx))
740                .collect(),
741        ),
742        serde_json::Value::Object(map) => serde_json::Value::Object(
743            map.iter()
744                .map(|(key, value)| (key.clone(), project_json_value(value, config, ctx)))
745                .collect(),
746        ),
747        other => other.clone(),
748    }
749}
750
751fn batch_child_tool_name(batch_args: &serde_json::Value, index: usize) -> Option<&str> {
752    batch_args
753        .get("tool_calls")
754        .and_then(|value| value.as_array())
755        .and_then(|items| items.get(index))
756        .and_then(|value| value.get("tool"))
757        .and_then(|value| value.as_str())
758}
759
760fn batch_child_args(batch_args: &serde_json::Value, index: usize) -> serde_json::Value {
761    batch_args
762        .get("tool_calls")
763        .and_then(|value| value.as_array())
764        .and_then(|items| items.get(index))
765        .and_then(|value| value.get("parameters"))
766        .cloned()
767        .unwrap_or_else(|| serde_json::Value::Object(Default::default()))
768}
769
770#[cfg(test)]
771mod tests {
772    use super::*;
773    use serde_json::json;
774
775    #[test]
776    fn windowed_truncation_truncates_over_long_single_line_instead_of_dropping_it() {
777        // A single line longer than the whole byte budget must be cut at a
778        // char boundary, not dropped (which would leave an empty preview).
779        let line = "x".repeat(1000);
780        let got = truncate_windowed(
781            &line,
782            &WindowedTruncation {
783                max_lines: 400,
784                max_bytes: 64,
785                direction: TruncationDirection::Head,
786                unit: TruncationUnit::Bytes,
787                hint: "hint",
788            },
789        );
790        let preview = got.split("\n\n...").next().expect("preview");
791        assert!(!preview.is_empty(), "preview must not be empty: {got:?}");
792        assert!(preview.len() <= 64);
793        assert!(preview.chars().all(|c| c == 'x'));
794        assert!(got.contains("bytes truncated"));
795    }
796
797    #[test]
798    fn windowed_truncation_never_splits_a_multibyte_char() {
799        // Budget lands mid-way through a 3-byte char; must back off to a
800        // boundary rather than panic or emit invalid UTF-8.
801        let line = "★".repeat(100); // each '★' is 3 bytes
802        let got = truncate_windowed(
803            &line,
804            &WindowedTruncation {
805                max_lines: 400,
806                max_bytes: 10, // not a multiple of 3
807                direction: TruncationDirection::Head,
808                unit: TruncationUnit::Bytes,
809                hint: "hint",
810            },
811        );
812        let preview = got.split("\n\n...").next().expect("preview");
813        assert!(!preview.is_empty());
814        assert!(preview.chars().all(|c| c == '★'));
815        assert_eq!(preview.len() % 3, 0, "must cut on a char boundary");
816        assert!(preview.len() <= 10);
817    }
818
819    #[test]
820    fn windowed_truncation_returns_input_unchanged_when_within_budget() {
821        let text = "a\nb\nc";
822        let got = truncate_windowed(
823            text,
824            &WindowedTruncation {
825                max_lines: 400,
826                max_bytes: 1024,
827                direction: TruncationDirection::Head,
828                unit: TruncationUnit::Bytes,
829                hint: "hint",
830            },
831        );
832        assert_eq!(got, text);
833    }
834
835    #[test]
836    fn truncates_strings_with_terminal_style_marker() {
837        let config = ToolOutputBudgetConfig {
838            mode: ToolOutputBudgetMode::Tokens,
839            limit: 5,
840            max_lines: DEFAULT_TOOL_OUTPUT_BUDGET_MAX_LINES,
841        };
842        let got = project_text(
843            "this is an example of a long output that should be truncated",
844            &config,
845            &ToolResultProjectionContext {
846                session_id: "root".to_string(),
847                call_id: "call".to_string(),
848                tool_name: "grep".to_string(),
849                args: json!({}),
850                output: lash_core::ToolCallOutput::success(json!("unused")),
851                duration_ms: 1,
852            },
853        );
854        assert!(got.contains("tokens truncated"));
855        assert!(got.contains("Full output saved to:"));
856    }
857
858    #[test]
859    fn truncation_hint_reuses_existing_full_output_path() {
860        let config = ToolOutputBudgetConfig {
861            limit: 512,
862            ..ToolOutputBudgetConfig::default()
863        };
864        let projected = project_tool_result(
865            &config,
866            ToolResultProjectionContext {
867                session_id: "root".to_string(),
868                call_id: "call".to_string(),
869                tool_name: "exec_command".to_string(),
870                args: json!({}),
871                output: lash_core::ToolCallOutput::success(json!({
872                    "output": "x".repeat(20_000),
873                    "full_output_path": "/tmp/existing-shell-output.log",
874                })),
875                duration_ms: 1,
876            },
877        );
878        let output = render_model_return_parts(&projected.parts);
879        assert!(output.contains("Full output saved to: /tmp/existing-shell-output.log"));
880    }
881
882    #[test]
883    fn model_projection_can_collapse_large_structured_payload_to_string() {
884        let config = ToolOutputBudgetConfig {
885            mode: ToolOutputBudgetMode::Bytes,
886            limit: 40,
887            max_lines: DEFAULT_TOOL_OUTPUT_BUDGET_MAX_LINES,
888        };
889        let projected = project_tool_result(
890            &config,
891            ToolResultProjectionContext {
892                session_id: "root".to_string(),
893                call_id: "call".to_string(),
894                tool_name: "search_tools".to_string(),
895                args: json!({}),
896                output: lash_core::ToolCallOutput::success(json!({
897                    "results": [{"output": "x".repeat(200)}]
898                })),
899                duration_ms: 1,
900            },
901        );
902        assert!(render_model_return_parts(&projected.parts).contains("bytes truncated"));
903    }
904
905    #[test]
906    fn batch_model_projection_preserves_projected_child_payloads() {
907        let projected = project_tool_result(
908            &ToolOutputBudgetConfig::default(),
909            ToolResultProjectionContext {
910                session_id: "root".to_string(),
911                call_id: "call".to_string(),
912                tool_name: "batch".to_string(),
913                args: json!({}),
914                output: lash_core::ToolCallOutput::success(json!({
915                    "results": [
916                        {"tool": "read_file", "success": true, "duration_ms": 1, "result": "very long child payload"},
917                        {"tool": "grep", "success": false, "duration_ms": 1, "error": "boom"}
918                    ]
919                })),
920                duration_ms: 1,
921            },
922        );
923        let projected_value: serde_json::Value =
924            serde_json::from_str(&render_model_return_parts(&projected.parts)).unwrap();
925        let results = projected_value
926            .get("results")
927            .and_then(|value| value.as_array())
928            .expect("results");
929        assert_eq!(results.len(), 2);
930        assert_eq!(
931            results[0].get("result"),
932            Some(&json!("very long child payload"))
933        );
934        assert_eq!(results[1].get("error"), Some(&json!("boom")));
935    }
936
937    #[test]
938    fn batch_history_projection_recursively_projects_child_payloads() {
939        let projected = project_tool_result(
940            &ToolOutputBudgetConfig {
941                limit: 8,
942                ..ToolOutputBudgetConfig::default()
943            },
944            ToolResultProjectionContext {
945                session_id: "root".to_string(),
946                call_id: "call".to_string(),
947                tool_name: "batch".to_string(),
948                args: json!({}),
949                output: lash_core::ToolCallOutput::success(json!({
950                    "results": [
951                        {"tool": "read_file", "success": true, "duration_ms": 1, "result": "child payload"},
952                        {"tool": "grep", "success": false, "duration_ms": 1, "error": "boom"}
953                    ]
954                })),
955                duration_ms: 1,
956            },
957        );
958        let projected_value: serde_json::Value =
959            serde_json::from_str(&render_model_return_parts(&projected.parts)).unwrap();
960        let details = projected_value
961            .get("results")
962            .and_then(|value| value.as_array())
963            .expect("results");
964        assert_eq!(details.len(), 2);
965        let child_result = details[0]
966            .get("result")
967            .and_then(|value| value.as_str())
968            .unwrap_or_default();
969        assert!(child_result.contains("truncated"));
970        assert_eq!(details[1].get("error"), Some(&json!("boom")));
971    }
972}