Skip to main content

traitclaw_core/traits/
output_processor.rs

1//! Tool output processing — pluggable post-processing of tool results.
2//!
3//! By default, [`TruncateProcessor`] limits tool output to 10,000 characters.
4
5/// Trait for processing tool output before adding it to the message context.
6///
7/// Implementations MUST be sync (no I/O) and fast.
8pub trait OutputProcessor: Send + Sync {
9    /// Process a tool output string and return the (possibly modified) result.
10    fn process(&self, output: String) -> String;
11}
12
13/// Returns output unchanged.
14pub struct NoopProcessor;
15
16impl OutputProcessor for NoopProcessor {
17    fn process(&self, output: String) -> String {
18        output
19    }
20}
21
22/// Truncates output exceeding `max_chars` with a `"...\n[output truncated]"` suffix.
23pub struct TruncateProcessor {
24    /// Maximum number of characters before truncation.
25    max_chars: usize,
26}
27
28impl Default for TruncateProcessor {
29    fn default() -> Self {
30        Self { max_chars: 10_000 }
31    }
32}
33
34impl TruncateProcessor {
35    /// Create a processor with a custom character limit.
36    #[must_use]
37    pub fn new(max_chars: usize) -> Self {
38        Self { max_chars }
39    }
40}
41
42impl OutputProcessor for TruncateProcessor {
43    fn process(&self, output: String) -> String {
44        // Count characters, not bytes, to avoid panicking on multi-byte UTF-8.
45        let char_count = output.chars().count();
46        if char_count <= self.max_chars {
47            output
48        } else {
49            // Find the byte offset of the Nth character boundary.
50            let byte_offset = output
51                .char_indices()
52                .nth(self.max_chars)
53                .map_or(output.len(), |(idx, _)| idx);
54            let mut truncated = output[..byte_offset].to_string();
55            truncated.push_str("...\n[output truncated]");
56            truncated
57        }
58    }
59}
60
61/// Composes multiple processors in order — output flows through each stage.
62pub struct ChainProcessor {
63    processors: Vec<Box<dyn OutputProcessor>>,
64}
65
66impl ChainProcessor {
67    /// Create a chain from a list of processors.
68    #[must_use]
69    pub fn new(processors: Vec<Box<dyn OutputProcessor>>) -> Self {
70        Self { processors }
71    }
72}
73
74impl OutputProcessor for ChainProcessor {
75    fn process(&self, mut output: String) -> String {
76        for p in &self.processors {
77            output = p.process(output);
78        }
79        output
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86
87    #[test]
88    fn test_noop_returns_unchanged() {
89        let p = NoopProcessor;
90        let input = "hello world".to_string();
91        assert_eq!(p.process(input.clone()), input);
92    }
93
94    #[test]
95    fn test_truncate_at_boundary() {
96        let p = TruncateProcessor::new(10);
97        let short = "12345".to_string();
98        assert_eq!(
99            p.process(short.clone()),
100            short,
101            "should not truncate short input"
102        );
103
104        let exact = "1234567890".to_string();
105        assert_eq!(
106            p.process(exact.clone()),
107            exact,
108            "should not truncate exact-length input"
109        );
110
111        let long = "12345678901".to_string();
112        let result = p.process(long);
113        assert!(
114            result.ends_with("[output truncated]"),
115            "should truncate: {result}"
116        );
117        assert!(
118            result.starts_with("1234567890"),
119            "should keep first 10 chars"
120        );
121    }
122
123    #[test]
124    fn test_chain_applies_in_order() {
125        struct UpperCase;
126        impl OutputProcessor for UpperCase {
127            fn process(&self, output: String) -> String {
128                output.to_uppercase()
129            }
130        }
131
132        let chain = ChainProcessor::new(vec![
133            Box::new(UpperCase),
134            Box::new(TruncateProcessor::new(5)),
135        ]);
136
137        let result = chain.process("hello world".to_string());
138        // First: "HELLO WORLD", then truncate to 5
139        assert!(result.starts_with("HELLO"), "got: {result}");
140        assert!(result.contains("[output truncated]"));
141    }
142}