lean_ctx/proxy/
tool_output.rs

1use serde_json::Value;
2
3use super::compress::compress_tool_result;
4use super::tool_kind::{ToolResultKind, should_protect};
5
6enum JsonRewrite {
7    NotJson,
8    Unchanged,
9    Changed(String),
10}
11
12pub(super) fn compress_text(
13    text: &mut String,
14    tool_name: Option<&str>,
15    kind: ToolResultKind,
16) -> bool {
17    match rewrite_json_payload_text(text, kind, |inner| {
18        if should_protect(kind, inner) {
19            return None;
20        }
21        let compressed = compress_tool_result(inner, tool_name);
22        (compressed.len() < inner.len()).then_some(compressed)
23    }) {
24        JsonRewrite::Changed(compressed) => {
25            *text = compressed;
26            return true;
27        }
28        JsonRewrite::Unchanged => return false,
29        JsonRewrite::NotJson => {}
30    }
31
32    if should_protect(kind, text) {
33        return false;
34    }
35    let compressed = compress_tool_result(text, tool_name);
36    if compressed.len() < text.len() {
37        *text = compressed;
38        return true;
39    }
40    false
41}
42
43pub(super) fn compress_value(
44    value: &mut Value,
45    tool_name: Option<&str>,
46    kind: ToolResultKind,
47) -> bool {
48    match value {
49        Value::String(text) => compress_text(text, tool_name, kind),
50        Value::Array(parts) => {
51            let mut changed = false;
52            for part in parts.iter_mut() {
53                if let Some(Value::String(text)) = part.get_mut("text") {
54                    changed |= compress_text(text, tool_name, kind);
55                }
56            }
57            changed
58        }
59        _ => false,
60    }
61}
62
63pub(super) fn prune_text(text: &mut String, kind: ToolResultKind) -> bool {
64    match rewrite_json_payload_text(text, kind, |inner| {
65        super::history_prune::prune_output_text(inner, kind)
66    }) {
67        JsonRewrite::Changed(pruned) => {
68            *text = pruned;
69            return true;
70        }
71        JsonRewrite::Unchanged => return false,
72        JsonRewrite::NotJson => {}
73    }
74
75    if let Some(pruned) = super::history_prune::prune_output_text(text, kind) {
76        *text = pruned;
77        return true;
78    }
79    false
80}
81
82pub(super) fn prune_value(value: &mut Value, kind: ToolResultKind) -> bool {
83    match value {
84        Value::String(text) => prune_text(text, kind),
85        Value::Array(parts) => {
86            let mut changed = false;
87            for part in parts.iter_mut() {
88                if let Some(Value::String(text)) = part.get_mut("text") {
89                    changed |= prune_text(text, kind);
90                }
91            }
92            changed
93        }
94        _ => false,
95    }
96}
97
98fn rewrite_json_payload_text(
99    text: &str,
100    kind: ToolResultKind,
101    mut rewrite: impl FnMut(&str) -> Option<String>,
102) -> JsonRewrite {
103    let trimmed = text.trim();
104    if !(trimmed.starts_with('{') || trimmed.starts_with('[')) {
105        return JsonRewrite::NotJson;
106    }
107    let Ok(mut value) = serde_json::from_str::<Value>(trimmed) else {
108        return JsonRewrite::NotJson;
109    };
110    let mut touched = false;
111    let mut changed = false;
112    rewrite_json_text_values(&mut value, kind, &mut rewrite, &mut touched, &mut changed);
113    if !touched || !changed {
114        return JsonRewrite::Unchanged;
115    }
116    match serde_json::to_string(&value) {
117        Ok(serialized) if serialized.len() < text.len() => JsonRewrite::Changed(serialized),
118        _ => JsonRewrite::Unchanged,
119    }
120}
121
122fn rewrite_json_text_values(
123    value: &mut Value,
124    kind: ToolResultKind,
125    rewrite: &mut impl FnMut(&str) -> Option<String>,
126    touched: &mut bool,
127    changed: &mut bool,
128) {
129    match value {
130        Value::Object(map) => {
131            let is_text_part = map
132                .get("type")
133                .and_then(Value::as_str)
134                .is_some_and(|t| matches!(t, "text" | "input_text" | "output_text"));
135            let rewrite_all_strings =
136                matches!(kind, ToolResultKind::Shell | ToolResultKind::Search);
137            for (key, child) in map.iter_mut() {
138                if let Value::String(s) = child
139                    && (rewrite_all_strings || (is_text_part && key == "text"))
140                {
141                    *touched = true;
142                    if let Some(next) = rewrite(s) {
143                        *s = next;
144                        *changed = true;
145                    }
146                    continue;
147                }
148                rewrite_json_text_values(child, kind, rewrite, touched, changed);
149            }
150        }
151        Value::Array(items) => {
152            for item in items {
153                rewrite_json_text_values(item, kind, rewrite, touched, changed);
154            }
155        }
156        _ => {}
157    }
158}
lean_ctx/proxy/tool_output.rs

lean_ctx/proxy/
tool_output.rs
