llm_tool/types/truncation/
policy.rs

1//! Truncation policy definition and core middle-out execution logic.
2
3use alloc::{
4    borrow::Cow,
5    format,
6    string::{String, ToString},
7    sync::Arc,
8};
9
10use super::{
11    alignment::{MAX_LINE_SCAN_WINDOW, TruncationAlignment},
12    outcome::TruncationOutcome,
13    placeholder::MarkerPlaceholder,
14    ratio::HeadRatio,
15    sink::{OutputSpillSink, SpillContext, SpillDisposition, SpillError, SpillRef},
16};
17use crate::types::ToolOutput;
18
19/// Prefix of the in-band warning emitted when a configured [`OutputSpillSink`] fails.
20pub const SPILL_FAILURE_NOTICE_PREFIX: &str = " WARNING: full output could NOT be retained (";
21
22/// Suffix of the in-band warning emitted when a configured [`OutputSpillSink`] fails.
23pub const SPILL_FAILURE_NOTICE_SUFFIX: &str = "); the omitted bytes are unrecoverable.";
24
25/// Marker used when `max_bytes` is too small to hold the fully informative marker.
26///
27/// Truncation is never allowed to be silent, so content is sacrificed before the
28/// notice is: a reader must always be able to tell that bytes are missing.
29pub const MINIMAL_TRUNCATION_MARKER: &str = "\n[...truncated...]\n";
30
31/// Minimal marker variant used when a configured [`OutputSpillSink`] failed.
32pub const MINIMAL_TRUNCATION_MARKER_SPILL_FAILED: &str = "\n[...truncated; NOT retained...]\n";
33
34/// Render the human- and model-readable notice describing a failed spill.
35fn spill_failure_notice(error: &SpillError) -> String {
36    format!("{SPILL_FAILURE_NOTICE_PREFIX}{error}{SPILL_FAILURE_NOTICE_SUFFIX}")
37}
38
39/// Configuration policy for middle-out string truncation.
40#[derive(Debug, Clone)]
41pub struct TruncationPolicy {
42    /// Maximum allowed bytes in the truncated output.
43    pub max_bytes: usize,
44    /// Ratio of available content budget assigned to the head portion.
45    /// Default is 50/50 balance between head and tail.
46    pub head_ratio: HeadRatio,
47    /// Alignment boundary strategy.
48    pub alignment: TruncationAlignment,
49    /// Optional custom marker template.
50    pub marker_template: Option<String>,
51    /// Optional spill sink for retaining full un-truncated observations.
52    pub spill_sink: Option<Arc<dyn OutputSpillSink>>,
53}
54
55impl PartialEq for TruncationPolicy {
56    fn eq(&self, other: &Self) -> bool {
57        self.max_bytes == other.max_bytes
58            && self.head_ratio == other.head_ratio
59            && self.alignment == other.alignment
60            && self.marker_template == other.marker_template
61            && match (&self.spill_sink, &other.spill_sink) {
62                (None, None) => true,
63                (Some(a), Some(b)) => Arc::ptr_eq(a, b),
64                _ => false,
65            }
66    }
67}
68
69impl Eq for TruncationPolicy {}
70
71impl TruncationPolicy {
72    /// Create a new truncation policy with a specified `max_bytes` ceiling,
73    /// balanced 50/50 head-tail ratio, and character boundary alignment.
74    #[must_use]
75    pub const fn new(max_bytes: usize) -> Self {
76        Self {
77            max_bytes,
78            head_ratio: HeadRatio::BALANCED,
79            alignment: TruncationAlignment::CharBoundary,
80            marker_template: None,
81            spill_sink: None,
82        }
83    }
84
85    /// Set the head budget ratio.
86    #[must_use]
87    pub fn with_head_ratio(mut self, ratio: impl Into<HeadRatio>) -> Self {
88        self.head_ratio = ratio.into();
89        self
90    }
91
92    /// Set the boundary alignment strategy.
93    #[must_use]
94    pub const fn with_alignment(mut self, alignment: TruncationAlignment) -> Self {
95        self.alignment = alignment;
96        self
97    }
98
99    /// Set a custom marker template.
100    ///
101    /// Placeholders `{omitted_bytes}`, `{original_bytes}`, `{offset_start}`,
102    /// `{offset_end}`, and `{ref}` will be substituted. Legacy aliases `{count}`,
103    /// `{bytes}`, and `<N>` are also supported.
104    #[must_use]
105    pub fn with_marker(mut self, marker: impl Into<String>) -> Self {
106        self.marker_template = Some(marker.into());
107        self
108    }
109
110    /// Set a spill sink for retaining un-truncated observations.
111    #[must_use]
112    pub fn with_spill_sink(mut self, sink: Arc<dyn OutputSpillSink>) -> Self {
113        self.spill_sink = Some(sink);
114        self
115    }
116
117    /// Format the marker string for the given omitted bytes without sink reference.
118    ///
119    /// Preserved for backwards compatibility with existing callers.
120    #[must_use]
121    pub fn format_marker(&self, omitted_bytes: usize) -> String {
122        self.format_marker_detailed(
123            omitted_bytes,
124            omitted_bytes,
125            0,
126            omitted_bytes,
127            &SpillDisposition::NotConfigured,
128        )
129    }
130
131    /// Format marker with full metadata substitutions including offsets and spill disposition.
132    ///
133    /// When `spill` is [`SpillDisposition::Failed`], the rendered marker always carries
134    /// an in-band warning that the omitted bytes are unrecoverable — including for custom
135    /// templates that do not mention the failure themselves — so retention loss can never
136    /// pass unnoticed by the model reading the output.
137    #[must_use]
138    pub fn format_marker_detailed(
139        &self,
140        omitted_bytes: usize,
141        original_bytes: usize,
142        offset_start: usize,
143        offset_end: usize,
144        spill: &SpillDisposition,
145    ) -> String {
146        let omitted_str = omitted_bytes.to_string();
147        let mut marker = match &self.marker_template {
148            Some(template) => {
149                let orig_str = original_bytes.to_string();
150                let start_str = offset_start.to_string();
151                let end_str = offset_end.to_string();
152                let ref_str = spill.spill_ref().map_or("", SpillRef::as_str);
153
154                template
155                    .replace(MarkerPlaceholder::TOKEN_OMITTED_BYTES, &omitted_str)
156                    .replace(MarkerPlaceholder::TOKEN_ORIGINAL_BYTES, &orig_str)
157                    .replace(MarkerPlaceholder::TOKEN_OFFSET_START, &start_str)
158                    .replace(MarkerPlaceholder::TOKEN_OFFSET_END, &end_str)
159                    .replace(MarkerPlaceholder::TOKEN_SPILL_REF, ref_str)
160                    .replace(MarkerPlaceholder::DEPRECATED_COUNT, &omitted_str)
161                    .replace(MarkerPlaceholder::DEPRECATED_BYTES, &omitted_str)
162                    .replace(MarkerPlaceholder::DEPRECATED_N, &omitted_str)
163            }
164            None => match spill {
165                SpillDisposition::Spilled(handle) => format!(
166                    "\n[... {omitted_str} bytes omitted (bytes {offset_start}..{offset_end} of {original_bytes}). Full output retained as ref=\"{handle}\" — recall with read_observation. ...]\n"
167                ),
168                SpillDisposition::Failed(error) => format!(
169                    "\n[... {omitted_str} bytes omitted (bytes {offset_start}..{offset_end} of {original_bytes}).{notice} ...]\n",
170                    notice = spill_failure_notice(error)
171                ),
172                SpillDisposition::NotConfigured => {
173                    format!("\n[... {omitted_str} bytes truncated ...]\n")
174                }
175            },
176        };
177
178        // A custom template cannot be trusted to mention the failure, so append it.
179        if let Some(error) = spill.error()
180            && self.marker_template.is_some()
181        {
182            marker.push_str(&spill_failure_notice(error));
183        }
184
185        marker
186    }
187
188    /// Truncate `input` using this policy, returning a borrowed `Cow` when no truncation
189    /// is needed, or an owned string when middle-out truncation occurs.
190    ///
191    /// If a configured [`OutputSpillSink`] fails, this method does **not** fail: the
192    /// error is logged and rendered into the truncation marker, so the loss of the
193    /// omitted bytes is visible to both the operator and the model. Callers that need
194    /// to handle spill failure programmatically should use
195    /// [`TruncationPolicy::truncate_detailed`], which returns the [`SpillError`].
196    #[must_use]
197    pub fn truncate<'a>(&self, input: &'a str) -> Cow<'a, str> {
198        self.truncate_lossy_with_context(input, &SpillContext::ANONYMOUS)
199            .text
200    }
201
202    /// Infallible truncation that degrades visibly — rather than failing — when a
203    /// configured [`OutputSpillSink`] returns an error.
204    #[must_use]
205    pub fn truncate_lossy_with_context<'a>(
206        &self,
207        input: &'a str,
208        ctx: &SpillContext<'_>,
209    ) -> TruncationOutcome<'a> {
210        if let Some(outcome) = Self::passthrough(input, self.max_bytes) {
211            return outcome;
212        }
213
214        let spill = self.resolve_spill(ctx, input);
215        if let Some(error) = spill.error() {
216            tracing::error!(
217                error = %error,
218                tool = ctx.tool_name,
219                "spill sink failed; omitted bytes are unrecoverable and the truncation marker will say so"
220            );
221        }
222        self.truncate_oversized(input, &spill)
223    }
224
225    /// Truncate `input` using an anonymous spill context, returning a detailed [`TruncationOutcome`].
226    ///
227    /// # Errors
228    ///
229    /// Returns [`SpillError`] if an [`OutputSpillSink`] is configured and fails.
230    pub fn truncate_detailed<'a>(
231        &self,
232        input: &'a str,
233    ) -> Result<TruncationOutcome<'a>, SpillError> {
234        self.truncate_detailed_with_context(input, &SpillContext::ANONYMOUS)
235    }
236
237    /// Truncate `input` with caller-provided [`SpillContext`], returning a detailed [`TruncationOutcome`].
238    ///
239    /// # Errors
240    ///
241    /// Returns [`SpillError`] if an [`OutputSpillSink`] is configured and fails.
242    pub fn truncate_detailed_with_context<'a>(
243        &self,
244        input: &'a str,
245        ctx: &SpillContext<'_>,
246    ) -> Result<TruncationOutcome<'a>, SpillError> {
247        match Self::passthrough(input, self.max_bytes) {
248            Some(outcome) => Ok(outcome),
249            None => match self.resolve_spill(ctx, input) {
250                SpillDisposition::Failed(error) => Err(error),
251                spill => Ok(self.truncate_oversized(input, &spill)),
252            },
253        }
254    }
255
256    /// Return the untouched input when it already fits within `max_bytes`.
257    ///
258    /// Checked before any spill is attempted: content that is not truncated loses
259    /// nothing, so it must not consume observation-store capacity.
260    fn passthrough(input: &str, max_bytes: usize) -> Option<TruncationOutcome<'_>> {
261        let total_bytes = input.len();
262        (total_bytes <= max_bytes).then_some(TruncationOutcome {
263            text: Cow::Borrowed(input),
264            original_len: total_bytes,
265            omitted: None,
266            spill_ref: None,
267        })
268    }
269
270    /// Attempt to retain the full output, classifying the result.
271    fn resolve_spill(&self, ctx: &SpillContext<'_>, input: &str) -> SpillDisposition {
272        match &self.spill_sink {
273            Some(sink) => match sink.spill(ctx, input) {
274                Ok(handle) => SpillDisposition::Spilled(handle),
275                Err(error) => SpillDisposition::Failed(error),
276            },
277            None => SpillDisposition::NotConfigured,
278        }
279    }
280
281    /// Perform middle-out elision on input already known to exceed `max_bytes`,
282    /// with the spill attempt already resolved.
283    fn truncate_oversized<'a>(
284        &self,
285        input: &'a str,
286        spill: &SpillDisposition,
287    ) -> TruncationOutcome<'a> {
288        let total_bytes = input.len();
289        let spill_ref = spill.spill_ref().cloned();
290
291        if self.max_bytes == 0 {
292            return TruncationOutcome {
293                text: Cow::Borrowed(""),
294                original_len: total_bytes,
295                omitted: Some(0..total_bytes),
296                spill_ref,
297            };
298        }
299
300        // Estimate marker length using approximate omitted byte count.
301        let est_omitted = total_bytes.saturating_sub(self.max_bytes);
302        let sample_marker =
303            self.format_marker_detailed(est_omitted, total_bytes, 0, est_omitted, spill);
304        let marker_len = sample_marker.len();
305
306        if self.max_bytes <= marker_len {
307            return self.truncate_with_fallback_marker(input, spill, spill_ref);
308        }
309
310        let content_budget = self.max_bytes.saturating_sub(marker_len);
311        let (head_end, tail_start) =
312            self.compute_initial_boundaries(input, total_bytes, content_budget);
313
314        // If boundaries overlap or touch, no truncation is necessary
315        if head_end >= tail_start {
316            return TruncationOutcome {
317                text: Cow::Borrowed(input),
318                original_len: total_bytes,
319                omitted: None,
320                spill_ref,
321            };
322        }
323
324        let (head_end, tail_start, marker, current_total) =
325            self.converge_boundaries(input, total_bytes, head_end, tail_start, spill);
326
327        let mut output = String::with_capacity(current_total);
328        output.push_str(&input[..head_end]);
329        output.push_str(&marker);
330        output.push_str(&input[tail_start..]);
331
332        if output.len() > self.max_bytes {
333            let mut cut = self.max_bytes;
334            while cut > 0 && !output.is_char_boundary(cut) {
335                cut -= 1;
336            }
337            output.truncate(cut);
338        }
339
340        TruncationOutcome {
341            text: Cow::Owned(output),
342            original_len: total_bytes,
343            omitted: Some(head_end..tail_start),
344            spill_ref,
345        }
346    }
347
348    /// Choose the most informative marker that fits inside `max_bytes`.
349    fn fallback_marker(spill: &SpillDisposition, max_bytes: usize) -> String {
350        if let SpillDisposition::Spilled(handle) = spill {
351            let with_ref = format!("\n[...truncated; ref=\"{handle}\"...]\n");
352            if with_ref.len() <= max_bytes {
353                return with_ref;
354            }
355        }
356        match spill {
357            SpillDisposition::Failed(_) => MINIMAL_TRUNCATION_MARKER_SPILL_FAILED.into(),
358            SpillDisposition::NotConfigured | SpillDisposition::Spilled(_) => {
359                MINIMAL_TRUNCATION_MARKER.into()
360            }
361        }
362    }
363
364    /// Truncate when `max_bytes` cannot hold the fully informative marker.
365    ///
366    /// Content is sacrificed before the marker is, so the reader can always tell
367    /// that bytes are missing; a budget too small even for the minimal marker
368    /// yields a truncated marker rather than a bare, deceptively complete-looking
369    /// slice of the input.
370    fn truncate_with_fallback_marker<'a>(
371        &self,
372        input: &str,
373        spill: &SpillDisposition,
374        spill_ref: Option<SpillRef>,
375    ) -> TruncationOutcome<'a> {
376        let total_bytes = input.len();
377        let marker = Self::fallback_marker(spill, self.max_bytes);
378
379        if marker.len() >= self.max_bytes {
380            let mut cut = self.max_bytes;
381            while cut > 0 && !marker.is_char_boundary(cut) {
382                cut -= 1;
383            }
384            return TruncationOutcome {
385                text: Cow::Owned(marker[..cut].into()),
386                original_len: total_bytes,
387                omitted: Some(0..total_bytes),
388                spill_ref,
389            };
390        }
391
392        let mut head_end = self.max_bytes - marker.len();
393        while head_end > 0 && !input.is_char_boundary(head_end) {
394            head_end -= 1;
395        }
396
397        let mut text = String::with_capacity(head_end + marker.len());
398        text.push_str(&input[..head_end]);
399        text.push_str(&marker);
400
401        TruncationOutcome {
402            text: Cow::Owned(text),
403            original_len: total_bytes,
404            omitted: Some(head_end..total_bytes),
405            spill_ref,
406        }
407    }
408
409    fn compute_initial_boundaries(
410        &self,
411        input: &str,
412        total_bytes: usize,
413        content_budget: usize,
414    ) -> (usize, usize) {
415        let head_budget = self.head_ratio.compute_head_budget(content_budget);
416        let tail_budget = content_budget.saturating_sub(head_budget);
417
418        // Floor to valid char boundary for head
419        let mut head_end = head_budget.min(total_bytes);
420        while head_end > 0 && !input.is_char_boundary(head_end) {
421            head_end -= 1;
422        }
423
424        // Ceil to valid char boundary for tail
425        let tail_target = total_bytes.saturating_sub(tail_budget);
426        let mut tail_start = tail_target.min(total_bytes);
427        while tail_start < total_bytes && !input.is_char_boundary(tail_start) {
428            tail_start += 1;
429        }
430
431        // Apply line boundary alignment if configured
432        if self.alignment == TruncationAlignment::LineBoundary {
433            let head_scan_start = head_end.saturating_sub(MAX_LINE_SCAN_WINDOW);
434            if let Some(rel_nl) = input[head_scan_start..head_end].rfind('\n') {
435                head_end = head_scan_start + rel_nl + 1;
436            }
437
438            let tail_scan_end = (tail_start + MAX_LINE_SCAN_WINDOW).min(total_bytes);
439            if let Some(rel_nl) = input[tail_start..tail_scan_end].find('\n') {
440                tail_start = tail_start + rel_nl + 1;
441            }
442        }
443
444        (head_end, tail_start)
445    }
446
447    fn shrink_head_boundary(&self, input: &str, mut head_end: usize) -> usize {
448        if self.alignment == TruncationAlignment::LineBoundary {
449            let scan_start = head_end.saturating_sub(MAX_LINE_SCAN_WINDOW);
450            if let Some(rel_nl) = input[scan_start..head_end.saturating_sub(1)].rfind('\n') {
451                return scan_start + rel_nl + 1;
452            }
453        }
454        head_end = head_end.saturating_sub(1);
455        while head_end > 0 && !input.is_char_boundary(head_end) {
456            head_end -= 1;
457        }
458        head_end
459    }
460
461    fn expand_tail_boundary(
462        &self,
463        input: &str,
464        total_bytes: usize,
465        mut tail_start: usize,
466    ) -> usize {
467        if self.alignment == TruncationAlignment::LineBoundary {
468            let scan_start = tail_start.saturating_add(1).min(total_bytes);
469            let scan_end = (tail_start + MAX_LINE_SCAN_WINDOW).min(total_bytes);
470            if scan_start < scan_end
471                && let Some(rel_nl) = input[scan_start..scan_end].find('\n')
472            {
473                return scan_start + rel_nl + 1;
474            }
475        }
476        tail_start = tail_start.saturating_add(1);
477        while tail_start < total_bytes && !input.is_char_boundary(tail_start) {
478            tail_start += 1;
479        }
480        tail_start
481    }
482
483    fn converge_boundaries(
484        &self,
485        input: &str,
486        total_bytes: usize,
487        mut head_end: usize,
488        mut tail_start: usize,
489        spill: &SpillDisposition,
490    ) -> (usize, usize, String, usize) {
491        let mut omitted_bytes = tail_start.saturating_sub(head_end);
492        let mut marker =
493            self.format_marker_detailed(omitted_bytes, total_bytes, head_end, tail_start, spill);
494        let mut current_total = head_end + marker.len() + (total_bytes - tail_start);
495
496        while current_total > self.max_bytes && (head_end > 0 || tail_start < total_bytes) {
497            if head_end > 0 && (tail_start == total_bytes || head_end >= (total_bytes - tail_start))
498            {
499                head_end = self.shrink_head_boundary(input, head_end);
500            } else if tail_start < total_bytes {
501                tail_start = self.expand_tail_boundary(input, total_bytes, tail_start);
502            } else {
503                break;
504            }
505
506            omitted_bytes = tail_start.saturating_sub(head_end);
507            marker = self.format_marker_detailed(
508                omitted_bytes,
509                total_bytes,
510                head_end,
511                tail_start,
512                spill,
513            );
514            current_total = head_end + marker.len() + (total_bytes - tail_start);
515        }
516
517        (head_end, tail_start, marker, current_total)
518    }
519}
520
521/// Helper function to perform middle-out truncation on a string slice.
522#[must_use]
523pub fn truncate_middle<'a>(input: &'a str, policy: &TruncationPolicy) -> Cow<'a, str> {
524    policy.truncate(input)
525}
526
527/// Detailed middle-out truncation returning metadata and recall handle.
528///
529/// # Errors
530///
531/// Returns [`SpillError`] if an [`OutputSpillSink`] is configured and fails.
532pub fn truncate_middle_detailed<'a>(
533    input: &'a str,
534    policy: &TruncationPolicy,
535) -> Result<TruncationOutcome<'a>, SpillError> {
536    policy.truncate_detailed(input)
537}
538
539/// Detailed middle-out truncation with [`SpillContext`].
540///
541/// # Errors
542///
543/// Returns [`SpillError`] if an [`OutputSpillSink`] is configured and fails.
544pub fn truncate_middle_detailed_with_context<'a>(
545    input: &'a str,
546    policy: &TruncationPolicy,
547    ctx: &SpillContext<'_>,
548) -> Result<TruncationOutcome<'a>, SpillError> {
549    policy.truncate_detailed_with_context(input, ctx)
550}
551
552/// Helper function to truncate a [`ToolOutput`] in-place according to `policy`.
553pub fn truncate_output(output: &mut ToolOutput, policy: &TruncationPolicy) {
554    output.content = policy.truncate(&output.content).into_owned();
555}
556
557/// Detailed in-place truncation of [`ToolOutput`].
558///
559/// # Errors
560///
561/// Returns [`SpillError`] if an [`OutputSpillSink`] is configured and fails.
562pub fn truncate_output_detailed(
563    output: &mut ToolOutput,
564    policy: &TruncationPolicy,
565) -> Result<TruncationOutcome<'static>, SpillError> {
566    truncate_output_detailed_with_context(output, policy, &SpillContext::ANONYMOUS)
567}
568
569/// Detailed in-place truncation of [`ToolOutput`] with [`SpillContext`].
570///
571/// # Errors
572///
573/// Returns [`SpillError`] if an [`OutputSpillSink`] is configured and fails.
574pub fn truncate_output_detailed_with_context(
575    output: &mut ToolOutput,
576    policy: &TruncationPolicy,
577    ctx: &SpillContext<'_>,
578) -> Result<TruncationOutcome<'static>, SpillError> {
579    let outcome = policy.truncate_detailed_with_context(&output.content, ctx)?;
580    let text = outcome.text.into_owned();
581    let original_len = outcome.original_len;
582    let omitted = outcome.omitted;
583    let spill_ref = outcome.spill_ref;
584    output.content.clone_from(&text);
585    Ok(TruncationOutcome {
586        text: Cow::Owned(text),
587        original_len,
588        omitted,
589        spill_ref,
590    })
591}
llm_tool/types/truncation/policy.rs

llm_tool/types/truncation/
policy.rs
