lean_ctx/proxy/
google.rs

1use axum::{
2    body::Body,
3    extract::State,
4    http::{Request, StatusCode},
5    response::Response,
6};
7use serde_json::Value;
8
9use super::ProxyState;
10use super::compress::compress_tool_result;
11use super::forward;
12use super::tool_kind::{self, ToolResultKind, should_protect};
13
14pub async fn handler(
15    State(state): State<ProxyState>,
16    req: Request<Body>,
17) -> Result<Response, StatusCode> {
18    let upstream = state.gemini_upstream();
19    forward::forward_request(
20        State(state),
21        req,
22        &upstream,
23        "/",
24        compress_request_body,
25        "Gemini",
26        &["application/x-ndjson"],
27    )
28    .await
29}
30
31fn compress_request_body(parsed: Value, original_size: usize) -> (Vec<u8>, usize, usize) {
32    let mut doc = parsed;
33    let mut modified = false;
34
35    if let Some(contents) = doc.get_mut("contents").and_then(|c| c.as_array_mut()) {
36        // Gemini's implicit prompt cache is prefix-based, so the frozen OLD
37        // region is pruned at the same monotone staircase boundary as every
38        // other rail. We never remove a `contents` entry — only rewrite the
39        // `functionResponse` text in place — so the conversation structure
40        // (and any `functionCall` ↔ `functionResponse` correspondence) is intact.
41        let mode = crate::core::config::Config::load()
42            .proxy
43            .resolved_history_mode();
44        let boundary = super::history_prune::prune_boundary(mode, contents.len());
45
46        for (idx, content) in contents.iter_mut().enumerate() {
47            let in_old_region = idx < boundary;
48            let Some(parts) = content.get_mut("parts").and_then(|p| p.as_array_mut()) else {
49                continue;
50            };
51            for part in parts.iter_mut() {
52                let Some(func_resp) = part.get_mut("functionResponse") else {
53                    continue;
54                };
55                // Gemini carries the originating function name inline — route it
56                // to the compressor (not `None`) so tool-specific patterns apply.
57                let name = func_resp
58                    .get("name")
59                    .and_then(|v| v.as_str())
60                    .map(String::from);
61                let kind = name
62                    .as_deref()
63                    .map_or(ToolResultKind::Other, tool_kind::classify_tool_name);
64                let Some(response) = func_resp.get_mut("response") else {
65                    continue;
66                };
67                for field in ["result", "content"] {
68                    modified |= if in_old_region {
69                        prune_string_field(response, field, kind)
70                    } else {
71                        compress_string_field(response, field, name.as_deref(), kind)
72                    };
73                }
74            }
75        }
76    }
77
78    let out = serde_json::to_vec(&doc).unwrap_or_default();
79    let compressed_size = if modified { out.len() } else { original_size };
80    (out, original_size, compressed_size)
81}
82
83/// Compress a recent `functionResponse.response.<field>` string. `tool_name` is
84/// routed to the compressor so tool-specific patterns (git status, ls, …) apply;
85/// protected file/source reads in the recent region are left intact.
86fn compress_string_field(
87    obj: &mut Value,
88    field: &str,
89    tool_name: Option<&str>,
90    kind: ToolResultKind,
91) -> bool {
92    if let Some(val) = obj
93        .get_mut(field)
94        .and_then(|v| v.as_str().map(String::from))
95    {
96        if should_protect(kind, &val) {
97            return false;
98        }
99        let compressed = compress_tool_result(&val, tool_name);
100        if compressed.len() < val.len() {
101            obj[field] = Value::String(compressed);
102            return true;
103        }
104    }
105    false
106}
107
108/// Cache-aware prune of an OLD `functionResponse.response.<field>`: file/source
109/// reads collapse to a re-read stub, everything else head/tail summarizes.
110/// Content-deterministic, so the cached prefix stays byte-stable across turns.
111fn prune_string_field(obj: &mut Value, field: &str, kind: ToolResultKind) -> bool {
112    if let Some(val) = obj.get(field).and_then(|v| v.as_str())
113        && let Some(pruned) = super::history_prune::prune_output_text(val, kind)
114    {
115        obj[field] = Value::String(pruned);
116        return true;
117    }
118    false
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124
125    /// `pairs` Gemini turns: a `model` `functionCall` then the `user`
126    /// `functionResponse` carrying a long file read.
127    fn gemini_read_turns(pairs: usize) -> Vec<Value> {
128        let code = (0..40)
129            .map(|i| format!("    let v{i} = compute_{i}(ctx, opts);"))
130            .collect::<Vec<_>>()
131            .join("\n");
132        let mut contents = Vec::new();
133        for t in 0..pairs {
134            contents.push(serde_json::json!({
135                "role": "model",
136                "parts": [{"functionCall": {"name": "read_file", "args": {}}}]
137            }));
138            contents.push(serde_json::json!({
139                "role": "user",
140                "parts": [{"functionResponse": {
141                    "name": "read_file",
142                    "response": {"result": format!("{code}\n// turn {t}")}
143                }}]
144            }));
145        }
146        contents
147    }
148
149    #[test]
150    fn recent_response_routes_tool_name_to_compressor() {
151        // Default (isolated) config; single content → boundary 0 → recent path.
152        let _iso = crate::core::data_dir::isolated_data_dir();
153        // A compressible search result. The proxy must route the inline tool name
154        // to the shared engine, so its output matches the name-routed engine
155        // byte-for-byte (the contract that distinguishes this from `None`).
156        // `infer_command`'s use of the name is unit-tested in `compress.rs`.
157        let raw = (0..60)
158            .map(|i| format!("src/file_{i}.rs:{i}:    let matched = find(foo, bar, baz);"))
159            .collect::<Vec<_>>()
160            .join("\n");
161        let routed = compress_tool_result(&raw, Some("search_files"));
162        assert!(routed.len() < raw.len(), "fixture must be compressible");
163
164        let body = serde_json::json!({
165            "contents": [
166                {"role": "user", "parts": [{"functionResponse": {
167                    "name": "search_files", "response": {"result": raw}
168                }}]}
169            ]
170        });
171        let bytes = serde_json::to_vec(&body).unwrap();
172        let (out, orig, comp) = compress_request_body(body, bytes.len());
173        assert!(comp < orig, "recent response must be compressed");
174        let parsed: Value = serde_json::from_slice(&out).unwrap();
175        assert_eq!(
176            parsed["contents"][0]["parts"][0]["functionResponse"]["response"]["result"]
177                .as_str()
178                .unwrap(),
179            routed,
180            "Gemini path must route the inline tool name to the shared compressor"
181        );
182    }
183
184    #[test]
185    fn cache_aware_prune_stubs_old_reads_keeps_recent() {
186        let _iso = crate::core::data_dir::isolated_data_dir();
187        // 13 pairs = 26 contents → staircase boundary 16.
188        let contents = gemini_read_turns(13);
189        let n = contents.len();
190        let body = serde_json::json!({ "contents": contents });
191        let bytes = serde_json::to_vec(&body).unwrap();
192        let (out, orig, comp) = compress_request_body(body, bytes.len());
193        assert!(comp < orig, "old reads must be pruned for savings");
194
195        let parsed: Value = serde_json::from_slice(&out).unwrap();
196        let got = parsed["contents"].as_array().unwrap();
197        assert_eq!(got.len(), n, "no contents may be removed");
198        // OLD file read (content index 1, before boundary 16) is stubbed.
199        let old = got[1]["parts"][0]["functionResponse"]["response"]["result"]
200            .as_str()
201            .unwrap();
202        assert!(
203            old.contains("Re-read the file"),
204            "old read should be stubbed, got: {old}"
205        );
206        // RECENT file read (content index 25, after the boundary) keeps its body.
207        let recent = got[25]["parts"][0]["functionResponse"]["response"]["result"]
208            .as_str()
209            .unwrap();
210        assert!(
211            recent.contains("v39"),
212            "recent read must be protected, got: {recent}"
213        );
214    }
215
216    #[test]
217    fn short_history_is_passthrough() {
218        let _iso = crate::core::data_dir::isolated_data_dir();
219        let body = serde_json::json!({
220            "contents": [
221                {"role": "user", "parts": [{"functionResponse": {
222                    "name": "read_file", "response": {"result": "ok"}
223                }}]}
224            ]
225        });
226        let bytes = serde_json::to_vec(&body).unwrap();
227        let (out, orig, comp) = compress_request_body(body.clone(), bytes.len());
228        assert_eq!(comp, orig);
229        let reparsed: Value = serde_json::from_slice(&out).unwrap();
230        assert_eq!(reparsed, body);
231    }
232
233    #[test]
234    fn gemini_compression_is_deterministic() {
235        // #498: identical request → identical bytes.
236        let _iso = crate::core::data_dir::isolated_data_dir();
237        let mk = || serde_json::json!({ "contents": gemini_read_turns(13) });
238        let (a, b) = (mk(), mk());
239        let (la, lb) = (
240            serde_json::to_vec(&a).unwrap().len(),
241            serde_json::to_vec(&b).unwrap().len(),
242        );
243        let (out_a, _, _) = compress_request_body(a, la);
244        let (out_b, _, _) = compress_request_body(b, lb);
245        assert_eq!(out_a, out_b, "identical input must yield identical bytes");
246    }
247
248    #[test]
249    fn cache_aware_gemini_prefix_is_byte_stable_across_turns() {
250        // THE cache invariant for the Gemini rail: every `contents` entry before
251        // an already-passed boundary stays byte-identical as the chat grows.
252        let _iso = crate::core::data_dir::isolated_data_dir();
253        let mut prev: Vec<String> = Vec::new();
254        let mut prev_boundary = 0;
255        for pairs in 1..=20 {
256            let contents = gemini_read_turns(pairs);
257            let len = contents.len();
258            let body = serde_json::json!({ "contents": contents });
259            let bytes = serde_json::to_vec(&body).unwrap();
260            let (out, _, _) = compress_request_body(body, bytes.len());
261            let parsed: Value = serde_json::from_slice(&out).unwrap();
262            let items: Vec<String> = parsed["contents"]
263                .as_array()
264                .unwrap()
265                .iter()
266                .map(Value::to_string)
267                .collect();
268            for i in 0..prev_boundary {
269                assert_eq!(
270                    prev[i], items[i],
271                    "Gemini content {i} changed at turn {pairs} — prompt cache prefix broken"
272                );
273            }
274            prev = items;
275            prev_boundary = crate::proxy::history_prune::prune_boundary(
276                crate::core::config::HistoryMode::CacheAware,
277                len,
278            );
279        }
280    }
281}
lean_ctx/proxy/google.rs

lean_ctx/proxy/
google.rs
