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};
13use super::{cache_safety, prose};
14use crate::core::config::ProseRole;
15
16pub async fn handler(
17    State(state): State<ProxyState>,
18    req: Request<Body>,
19) -> Result<Response, StatusCode> {
20    let upstream = state.gemini_upstream();
21    forward::forward_request(
22        State(state),
23        req,
24        &upstream,
25        "/",
26        compress_request_body,
27        "Gemini",
28        &["application/x-ndjson"],
29    )
30    .await
31}
32
33fn compress_request_body(parsed: Value, original_size: usize) -> (Vec<u8>, usize, usize) {
34    let mut doc = parsed;
35    let mut modified = false;
36
37    // Opt-in per-role prose aggressiveness (#710); both default `None` → no-op.
38    let cfg = crate::core::config::Config::load();
39    let system_aggr = cfg.proxy.resolved_role_aggressiveness(ProseRole::System);
40    let user_aggr = cfg.proxy.resolved_role_aggressiveness(ProseRole::User);
41    let mut prose_segments: u64 = 0;
42
43    // System prose: the top-level `systemInstruction` anchor. Gemini has no
44    // client `cache_control`, and the rewrite is deterministic, so the implicit
45    // prefix cache stays byte-stable across turns — cache-safe by construction.
46    if let Some(a) = system_aggr {
47        for key in ["systemInstruction", "system_instruction"] {
48            if let Some(parts) = doc
49                .get_mut(key)
50                .and_then(|si| si.get_mut("parts"))
51                .and_then(|p| p.as_array_mut())
52            {
53                prose_segments += u64::from(prose::compress_gemini_text_parts(parts, a));
54            }
55        }
56    }
57
58    if let Some(contents) = doc.get_mut("contents").and_then(|c| c.as_array_mut()) {
59        // Gemini's implicit prompt cache is prefix-based, so the frozen OLD
60        // region is pruned at the same monotone staircase boundary as every
61        // other rail. We never remove a `contents` entry — only rewrite the
62        // `functionResponse` text in place — so the conversation structure
63        // (and any `functionCall` ↔ `functionResponse` correspondence) is intact.
64        let mode = cfg.proxy.resolved_history_mode();
65        let boundary = super::history_prune::prune_boundary(mode, contents.len());
66
67        for (idx, content) in contents.iter_mut().enumerate() {
68            let in_old_region = idx < boundary;
69            // Own the role before the mutable `parts` borrow below.
70            let role = content
71                .get("role")
72                .and_then(|r| r.as_str())
73                .map(String::from);
74            let Some(parts) = content.get_mut("parts").and_then(|p| p.as_array_mut()) else {
75                continue;
76            };
77            for part in parts.iter_mut() {
78                let Some(func_resp) = part.get_mut("functionResponse") else {
79                    continue;
80                };
81                // Gemini carries the originating function name inline — route it
82                // to the compressor (not `None`) so tool-specific patterns apply.
83                let name = func_resp
84                    .get("name")
85                    .and_then(|v| v.as_str())
86                    .map(String::from);
87                let kind = name
88                    .as_deref()
89                    .map_or(ToolResultKind::Other, tool_kind::classify_tool_name);
90                let Some(response) = func_resp.get_mut("response") else {
91                    continue;
92                };
93                for field in ["result", "content"] {
94                    modified |= if in_old_region {
95                        prune_string_field(response, field, kind)
96                    } else {
97                        compress_string_field(response, field, name.as_deref(), kind)
98                    };
99                }
100            }
101
102            // Frozen-region user prose: free-text `text` parts of user turns in
103            // the old region `[0, boundary)`. Model turns (assistant) and tool
104            // I/O parts are never touched.
105            if in_old_region
106                && role.as_deref() == Some("user")
107                && let Some(a) = user_aggr
108            {
109                prose_segments += u64::from(prose::compress_gemini_text_parts(parts, a));
110            }
111        }
112    }
113
114    if prose_segments > 0 {
115        modified = true;
116    }
117    cache_safety::record(prose_segments, true);
118
119    let out = serde_json::to_vec(&doc).unwrap_or_default();
120    let compressed_size = if modified { out.len() } else { original_size };
121    (out, original_size, compressed_size)
122}
123
124/// Compress a recent `functionResponse.response.<field>` string. `tool_name` is
125/// routed to the compressor so tool-specific patterns (git status, ls, …) apply;
126/// protected file/source reads in the recent region are left intact.
127fn compress_string_field(
128    obj: &mut Value,
129    field: &str,
130    tool_name: Option<&str>,
131    kind: ToolResultKind,
132) -> bool {
133    if let Some(val) = obj
134        .get_mut(field)
135        .and_then(|v| v.as_str().map(String::from))
136    {
137        if should_protect(kind, &val) {
138            return false;
139        }
140        let compressed = compress_tool_result(&val, tool_name);
141        if compressed.len() < val.len() {
142            obj[field] = Value::String(compressed);
143            return true;
144        }
145    }
146    false
147}
148
149/// Cache-aware prune of an OLD `functionResponse.response.<field>`: file/source
150/// reads collapse to a re-read stub, everything else head/tail summarizes.
151/// Content-deterministic, so the cached prefix stays byte-stable across turns.
152fn prune_string_field(obj: &mut Value, field: &str, kind: ToolResultKind) -> bool {
153    if let Some(val) = obj.get(field).and_then(|v| v.as_str())
154        && let Some(pruned) = super::history_prune::prune_output_text(val, kind)
155    {
156        obj[field] = Value::String(pruned);
157        return true;
158    }
159    false
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165
166    /// `pairs` Gemini turns: a `model` `functionCall` then the `user`
167    /// `functionResponse` carrying a long file read.
168    fn gemini_read_turns(pairs: usize) -> Vec<Value> {
169        let code = (0..40)
170            .map(|i| format!("    let v{i} = compute_{i}(ctx, opts);"))
171            .collect::<Vec<_>>()
172            .join("\n");
173        let mut contents = Vec::new();
174        for t in 0..pairs {
175            contents.push(serde_json::json!({
176                "role": "model",
177                "parts": [{"functionCall": {"name": "read_file", "args": {}}}]
178            }));
179            contents.push(serde_json::json!({
180                "role": "user",
181                "parts": [{"functionResponse": {
182                    "name": "read_file",
183                    "response": {"result": format!("{code}\n// turn {t}")}
184                }}]
185            }));
186        }
187        contents
188    }
189
190    #[test]
191    fn recent_response_routes_tool_name_to_compressor() {
192        // Default (isolated) config; single content → boundary 0 → recent path.
193        let _iso = crate::core::data_dir::isolated_data_dir();
194        // A compressible search result. The proxy must route the inline tool name
195        // to the shared engine, so its output matches the name-routed engine
196        // byte-for-byte (the contract that distinguishes this from `None`).
197        // `infer_command`'s use of the name is unit-tested in `compress.rs`.
198        let raw = (0..60)
199            .map(|i| format!("src/file_{i}.rs:{i}:    let matched = find(foo, bar, baz);"))
200            .collect::<Vec<_>>()
201            .join("\n");
202        let routed = compress_tool_result(&raw, Some("search_files"));
203        assert!(routed.len() < raw.len(), "fixture must be compressible");
204
205        let body = serde_json::json!({
206            "contents": [
207                {"role": "user", "parts": [{"functionResponse": {
208                    "name": "search_files", "response": {"result": raw}
209                }}]}
210            ]
211        });
212        let bytes = serde_json::to_vec(&body).unwrap();
213        let (out, orig, comp) = compress_request_body(body, bytes.len());
214        assert!(comp < orig, "recent response must be compressed");
215        let parsed: Value = serde_json::from_slice(&out).unwrap();
216        assert_eq!(
217            parsed["contents"][0]["parts"][0]["functionResponse"]["response"]["result"]
218                .as_str()
219                .unwrap(),
220            routed,
221            "Gemini path must route the inline tool name to the shared compressor"
222        );
223    }
224
225    #[test]
226    fn cache_aware_prune_stubs_old_reads_keeps_recent() {
227        let _iso = crate::core::data_dir::isolated_data_dir();
228        // 13 pairs = 26 contents → staircase boundary 16.
229        let contents = gemini_read_turns(13);
230        let n = contents.len();
231        let body = serde_json::json!({ "contents": contents });
232        let bytes = serde_json::to_vec(&body).unwrap();
233        let (out, orig, comp) = compress_request_body(body, bytes.len());
234        assert!(comp < orig, "old reads must be pruned for savings");
235
236        let parsed: Value = serde_json::from_slice(&out).unwrap();
237        let got = parsed["contents"].as_array().unwrap();
238        assert_eq!(got.len(), n, "no contents may be removed");
239        // OLD file read (content index 1, before boundary 16) is stubbed.
240        let old = got[1]["parts"][0]["functionResponse"]["response"]["result"]
241            .as_str()
242            .unwrap();
243        assert!(
244            old.contains("Re-read the file"),
245            "old read should be stubbed, got: {old}"
246        );
247        // RECENT file read (content index 25, after the boundary) keeps its body.
248        let recent = got[25]["parts"][0]["functionResponse"]["response"]["result"]
249            .as_str()
250            .unwrap();
251        assert!(
252            recent.contains("v39"),
253            "recent read must be protected, got: {recent}"
254        );
255    }
256
257    #[test]
258    fn short_history_is_passthrough() {
259        let _iso = crate::core::data_dir::isolated_data_dir();
260        let body = serde_json::json!({
261            "contents": [
262                {"role": "user", "parts": [{"functionResponse": {
263                    "name": "read_file", "response": {"result": "ok"}
264                }}]}
265            ]
266        });
267        let bytes = serde_json::to_vec(&body).unwrap();
268        let (out, orig, comp) = compress_request_body(body.clone(), bytes.len());
269        assert_eq!(comp, orig);
270        let reparsed: Value = serde_json::from_slice(&out).unwrap();
271        assert_eq!(reparsed, body);
272    }
273
274    #[test]
275    fn gemini_compression_is_deterministic() {
276        // #498: identical request → identical bytes.
277        let _iso = crate::core::data_dir::isolated_data_dir();
278        let mk = || serde_json::json!({ "contents": gemini_read_turns(13) });
279        let (a, b) = (mk(), mk());
280        let (la, lb) = (
281            serde_json::to_vec(&a).unwrap().len(),
282            serde_json::to_vec(&b).unwrap().len(),
283        );
284        let (out_a, _, _) = compress_request_body(a, la);
285        let (out_b, _, _) = compress_request_body(b, lb);
286        assert_eq!(out_a, out_b, "identical input must yield identical bytes");
287    }
288
289    fn big_prose() -> String {
290        let p = "You are a careful, senior software engineer. You always explain your \
291                 reasoning before making changes, you prefer small reviewable diffs, and \
292                 you never introduce mock data or placeholders into production code. ";
293        [p; 6].join("\n")
294    }
295
296    #[test]
297    fn system_instruction_compressed_and_model_untouched() {
298        let _iso = crate::core::data_dir::isolated_data_dir();
299        crate::core::config::Config::update_global(|c| {
300            c.proxy.role_aggressiveness.system = Some(0.6);
301        })
302        .unwrap();
303
304        let prose = big_prose();
305        let body = serde_json::json!({
306            "systemInstruction": {"parts": [{"text": prose}]},
307            "contents": [
308                {"role": "user", "parts": [{"text": "hi"}]},
309                {"role": "model", "parts": [{"text": prose}]},
310            ]
311        });
312        let bytes = serde_json::to_vec(&body).unwrap();
313        let (out, _o, _c) = compress_request_body(body, bytes.len());
314        let parsed: Value = serde_json::from_slice(&out).unwrap();
315
316        assert!(
317            parsed["systemInstruction"]["parts"][0]["text"]
318                .as_str()
319                .unwrap()
320                .len()
321                < prose.len(),
322            "systemInstruction prose must be compressed when enabled"
323        );
324        assert_eq!(
325            parsed["contents"][1]["parts"][0]["text"].as_str().unwrap(),
326            prose,
327            "model (assistant) turns must pass through verbatim (#710)"
328        );
329    }
330
331    #[test]
332    fn gemini_prose_compression_is_deterministic() {
333        let _iso = crate::core::data_dir::isolated_data_dir();
334        crate::core::config::Config::update_global(|c| {
335            c.proxy.role_aggressiveness.system = Some(0.6);
336        })
337        .unwrap();
338        let prose = big_prose();
339        let mk = || {
340            serde_json::json!({
341                "systemInstruction": {"parts": [{"text": prose}]},
342                "contents": [{"role": "user", "parts": [{"text": "hi"}]}]
343            })
344        };
345        let (a, b) = (mk(), mk());
346        let la = serde_json::to_vec(&a).unwrap().len();
347        let lb = serde_json::to_vec(&b).unwrap().len();
348        assert_eq!(
349            compress_request_body(a, la).0,
350            compress_request_body(b, lb).0,
351            "identical input must yield byte-identical output (#498)"
352        );
353    }
354
355    #[test]
356    fn cache_aware_gemini_prefix_is_byte_stable_across_turns() {
357        // THE cache invariant for the Gemini rail: every `contents` entry before
358        // an already-passed boundary stays byte-identical as the chat grows.
359        let _iso = crate::core::data_dir::isolated_data_dir();
360        let mut prev: Vec<String> = Vec::new();
361        let mut prev_boundary = 0;
362        for pairs in 1..=20 {
363            let contents = gemini_read_turns(pairs);
364            let len = contents.len();
365            let body = serde_json::json!({ "contents": contents });
366            let bytes = serde_json::to_vec(&body).unwrap();
367            let (out, _, _) = compress_request_body(body, bytes.len());
368            let parsed: Value = serde_json::from_slice(&out).unwrap();
369            let items: Vec<String> = parsed["contents"]
370                .as_array()
371                .unwrap()
372                .iter()
373                .map(Value::to_string)
374                .collect();
375            for i in 0..prev_boundary {
376                assert_eq!(
377                    prev[i], items[i],
378                    "Gemini content {i} changed at turn {pairs} — prompt cache prefix broken"
379                );
380            }
381            prev = items;
382            prev_boundary = crate::proxy::history_prune::prune_boundary(
383                crate::core::config::HistoryMode::CacheAware,
384                len,
385            );
386        }
387    }
388}
lean_ctx/proxy/google.rs

lean_ctx/proxy/
google.rs
