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::{HistoryMode, 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    // Gemini carries the model in the URL path, not the body — capture it here so
22    // the effort applier can pick the right thinking control (#840).
23    let model = super::usage::gemini_model_from_path(req.uri().path());
24    forward::forward_request(
25        State(state),
26        req,
27        &upstream,
28        "/",
29        move |body, size| compress_request_body(body, size, model.as_deref()),
30        "Gemini",
31        &["application/x-ndjson"],
32    )
33    .await
34}
35
36fn compress_request_body(
37    parsed: Value,
38    original_size: usize,
39    model: Option<&str>,
40) -> (Vec<u8>, usize, usize) {
41    let mut doc = parsed;
42    let mut modified = false;
43
44    // Opt-in per-role prose aggressiveness (#710); both default `None` → no-op.
45    let cfg = crate::core::config::Config::load();
46    let system_aggr = cfg.proxy.resolved_role_aggressiveness(ProseRole::System);
47    let user_aggr = cfg.proxy.resolved_role_aggressiveness(ProseRole::User);
48    let live_compress = cfg.proxy.live_compresses();
49    let mode = cfg.proxy.resolved_history_mode();
50    // #493: in-band CCR expansion (opt-in). Splice any <lc_expand:HASH> the model
51    // echoed back into the verbatim original from the local tee store. A strict
52    // no-op when no marker is present (byte-identical body → cache-safe). Runs
53    // before the meter-only short-circuit so an explicit expand request is
54    // honored even when the proxy is otherwise byte-passthrough.
55    if cfg.proxy.ccr_inband_enabled() {
56        modified |= super::ccr::splice_inband_in_place(&mut doc);
57    }
58    // #834/#840: cache-safe cross-provider effort control. Default off → no-op.
59    // The level is a constant, so it never perturbs the prompt-cache prefix; it
60    // sets generationConfig.thinkingConfig (thinkingLevel on 3.x, thinkingBudget
61    // on 2.5 pro/flash) only for models that accept it and only when the client
62    // didn't pin its own thinking field. `model` is read from the URL path.
63    if let Some(effort) = cfg.proxy.resolved_effort() {
64        modified |= super::effort::apply_google(&mut doc, effort, model);
65    }
66    // Meter-only (#481): no live compression, no history pruning, no prose → the
67    // body is forwarded unchanged while usage metering still runs. A pending
68    // in-band splice (`modified`) opts out: the body did change this turn.
69    if !live_compress
70        && mode == HistoryMode::Off
71        && system_aggr.is_none()
72        && user_aggr.is_none()
73        && !modified
74    {
75        let out = serde_json::to_vec(&doc).unwrap_or_default();
76        return (out, original_size, original_size);
77    }
78    let mut prose_segments: u64 = 0;
79
80    // System prose: the top-level `systemInstruction` anchor. Gemini has no
81    // client `cache_control`, and the rewrite is deterministic, so the implicit
82    // prefix cache stays byte-stable across turns — cache-safe by construction.
83    if let Some(a) = system_aggr {
84        for key in ["systemInstruction", "system_instruction"] {
85            if let Some(parts) = doc
86                .get_mut(key)
87                .and_then(|si| si.get_mut("parts"))
88                .and_then(|p| p.as_array_mut())
89            {
90                prose_segments += u64::from(prose::compress_gemini_text_parts(parts, a));
91            }
92        }
93    }
94
95    if let Some(contents) = doc.get_mut("contents").and_then(|c| c.as_array_mut()) {
96        // Gemini's implicit prompt cache is prefix-based, so the frozen OLD
97        // region is pruned at the same monotone staircase boundary as every
98        // other rail. We never remove a `contents` entry — only rewrite the
99        // `functionResponse` text in place — so the conversation structure
100        // (and any `functionCall` ↔ `functionResponse` correspondence) is intact.
101        // `mode` resolved above.
102        let boundary = super::history_prune::prune_boundary(mode, contents.len());
103
104        for (idx, content) in contents.iter_mut().enumerate() {
105            let in_old_region = idx < boundary;
106            // Own the role before the mutable `parts` borrow below.
107            let role = content
108                .get("role")
109                .and_then(|r| r.as_str())
110                .map(String::from);
111            let Some(parts) = content.get_mut("parts").and_then(|p| p.as_array_mut()) else {
112                continue;
113            };
114            for part in parts.iter_mut() {
115                let Some(func_resp) = part.get_mut("functionResponse") else {
116                    continue;
117                };
118                // Gemini carries the originating function name inline — route it
119                // to the compressor (not `None`) so tool-specific patterns apply.
120                let name = func_resp
121                    .get("name")
122                    .and_then(|v| v.as_str())
123                    .map(String::from);
124                let kind = name
125                    .as_deref()
126                    .map_or(ToolResultKind::Other, tool_kind::classify_tool_name);
127                // #481: recent-region live compression respects the global toggle
128                // and the per-tool exclusion list (Serena default). Old-region
129                // pruning stays governed by `history_mode`.
130                let live = live_compress
131                    && !name
132                        .as_deref()
133                        .is_some_and(|n| cfg.proxy.is_tool_live_compress_excluded(n));
134                let Some(response) = func_resp.get_mut("response") else {
135                    continue;
136                };
137                for field in ["result", "content"] {
138                    modified |= if in_old_region {
139                        prune_string_field(response, field, kind)
140                    } else if live {
141                        compress_string_field(response, field, name.as_deref(), kind)
142                    } else {
143                        false
144                    };
145                }
146            }
147
148            // Frozen-region user prose: free-text `text` parts of user turns in
149            // the old region `[0, boundary)`. Model turns (assistant) and tool
150            // I/O parts are never touched.
151            if in_old_region
152                && role.as_deref() == Some("user")
153                && let Some(a) = user_aggr
154            {
155                prose_segments += u64::from(prose::compress_gemini_text_parts(parts, a));
156            }
157        }
158    }
159
160    if prose_segments > 0 {
161        modified = true;
162    }
163    cache_safety::record(prose_segments, true);
164
165    let out = serde_json::to_vec(&doc).unwrap_or_default();
166    let compressed_size = if modified { out.len() } else { original_size };
167    (out, original_size, compressed_size)
168}
169
170/// Compress a recent `functionResponse.response.<field>` string. `tool_name` is
171/// routed to the compressor so tool-specific patterns (git status, ls, …) apply;
172/// protected file/source reads in the recent region are left intact.
173fn compress_string_field(
174    obj: &mut Value,
175    field: &str,
176    tool_name: Option<&str>,
177    kind: ToolResultKind,
178) -> bool {
179    if let Some(val) = obj
180        .get_mut(field)
181        .and_then(|v| v.as_str().map(String::from))
182    {
183        if should_protect(kind, &val) {
184            return false;
185        }
186        let compressed = compress_tool_result(&val, tool_name);
187        if compressed.len() < val.len() {
188            obj[field] = Value::String(compressed);
189            return true;
190        }
191    }
192    false
193}
194
195/// Cache-aware prune of an OLD `functionResponse.response.<field>`: file/source
196/// reads collapse to a re-read stub, everything else head/tail summarizes.
197/// Content-deterministic, so the cached prefix stays byte-stable across turns.
198fn prune_string_field(obj: &mut Value, field: &str, kind: ToolResultKind) -> bool {
199    if let Some(val) = obj.get(field).and_then(|v| v.as_str())
200        && let Some(pruned) = super::history_prune::prune_output_text(val, kind)
201    {
202        obj[field] = Value::String(pruned);
203        return true;
204    }
205    false
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211
212    /// `pairs` Gemini turns: a `model` `functionCall` then the `user`
213    /// `functionResponse` carrying a long file read.
214    fn gemini_read_turns(pairs: usize) -> Vec<Value> {
215        let code = (0..40)
216            .map(|i| format!("    let v{i} = compute_{i}(ctx, opts);"))
217            .collect::<Vec<_>>()
218            .join("\n");
219        let mut contents = Vec::new();
220        for t in 0..pairs {
221            contents.push(serde_json::json!({
222                "role": "model",
223                "parts": [{"functionCall": {"name": "read_file", "args": {}}}]
224            }));
225            contents.push(serde_json::json!({
226                "role": "user",
227                "parts": [{"functionResponse": {
228                    "name": "read_file",
229                    "response": {"result": format!("{code}\n// turn {t}")}
230                }}]
231            }));
232        }
233        contents
234    }
235
236    #[test]
237    fn recent_response_routes_tool_name_to_compressor() {
238        // Default (isolated) config; single content → boundary 0 → recent path.
239        let _iso = crate::core::data_dir::isolated_data_dir();
240        // A compressible search result. The proxy must route the inline tool name
241        // to the shared engine, so its output matches the name-routed engine
242        // byte-for-byte (the contract that distinguishes this from `None`).
243        // `infer_command`'s use of the name is unit-tested in `compress.rs`.
244        let raw = (0..60)
245            .map(|i| format!("src/file_{i}.rs:{i}:    let matched = find(foo, bar, baz);"))
246            .collect::<Vec<_>>()
247            .join("\n");
248        let routed = compress_tool_result(&raw, Some("search_files"));
249        assert!(routed.len() < raw.len(), "fixture must be compressible");
250
251        let body = serde_json::json!({
252            "contents": [
253                {"role": "user", "parts": [{"functionResponse": {
254                    "name": "search_files", "response": {"result": raw}
255                }}]}
256            ]
257        });
258        let bytes = serde_json::to_vec(&body).unwrap();
259        let (out, orig, comp) = compress_request_body(body, bytes.len(), None);
260        assert!(comp < orig, "recent response must be compressed");
261        let parsed: Value = serde_json::from_slice(&out).unwrap();
262        assert_eq!(
263            parsed["contents"][0]["parts"][0]["functionResponse"]["response"]["result"]
264                .as_str()
265                .unwrap(),
266            routed,
267            "Gemini path must route the inline tool name to the shared compressor"
268        );
269    }
270
271    #[test]
272    fn cache_aware_prune_stubs_old_reads_keeps_recent() {
273        let _iso = crate::core::data_dir::isolated_data_dir();
274        // 13 pairs = 26 contents → staircase boundary 16.
275        let contents = gemini_read_turns(13);
276        let n = contents.len();
277        let body = serde_json::json!({ "contents": contents });
278        let bytes = serde_json::to_vec(&body).unwrap();
279        let (out, orig, comp) = compress_request_body(body, bytes.len(), None);
280        assert!(comp < orig, "old reads must be pruned for savings");
281
282        let parsed: Value = serde_json::from_slice(&out).unwrap();
283        let got = parsed["contents"].as_array().unwrap();
284        assert_eq!(got.len(), n, "no contents may be removed");
285        // OLD file read (content index 1, before boundary 16) is stubbed.
286        let old = got[1]["parts"][0]["functionResponse"]["response"]["result"]
287            .as_str()
288            .unwrap();
289        assert!(
290            old.contains("Re-read the file"),
291            "old read should be stubbed, got: {old}"
292        );
293        // RECENT file read (content index 25, after the boundary) keeps its body.
294        let recent = got[25]["parts"][0]["functionResponse"]["response"]["result"]
295            .as_str()
296            .unwrap();
297        assert!(
298            recent.contains("v39"),
299            "recent read must be protected, got: {recent}"
300        );
301    }
302
303    #[test]
304    fn short_history_is_passthrough() {
305        let _iso = crate::core::data_dir::isolated_data_dir();
306        let body = serde_json::json!({
307            "contents": [
308                {"role": "user", "parts": [{"functionResponse": {
309                    "name": "read_file", "response": {"result": "ok"}
310                }}]}
311            ]
312        });
313        let bytes = serde_json::to_vec(&body).unwrap();
314        let (out, orig, comp) = compress_request_body(body.clone(), bytes.len(), None);
315        assert_eq!(comp, orig);
316        let reparsed: Value = serde_json::from_slice(&out).unwrap();
317        assert_eq!(reparsed, body);
318    }
319
320    #[test]
321    fn gemini_compression_is_deterministic() {
322        // #498: identical request → identical bytes.
323        let _iso = crate::core::data_dir::isolated_data_dir();
324        let mk = || serde_json::json!({ "contents": gemini_read_turns(13) });
325        let (a, b) = (mk(), mk());
326        let (la, lb) = (
327            serde_json::to_vec(&a).unwrap().len(),
328            serde_json::to_vec(&b).unwrap().len(),
329        );
330        let (out_a, _, _) = compress_request_body(a, la, None);
331        let (out_b, _, _) = compress_request_body(b, lb, None);
332        assert_eq!(out_a, out_b, "identical input must yield identical bytes");
333    }
334
335    fn big_prose() -> String {
336        let p = "You are a careful, senior software engineer. You always explain your \
337                 reasoning before making changes, you prefer small reviewable diffs, and \
338                 you never introduce mock data or placeholders into production code. ";
339        [p; 6].join("\n")
340    }
341
342    #[test]
343    fn system_instruction_compressed_and_model_untouched() {
344        let _iso = crate::core::data_dir::isolated_data_dir();
345        crate::core::config::Config::update_global(|c| {
346            c.proxy.role_aggressiveness.system = Some(0.6);
347        })
348        .unwrap();
349
350        let prose = big_prose();
351        let body = serde_json::json!({
352            "systemInstruction": {"parts": [{"text": prose}]},
353            "contents": [
354                {"role": "user", "parts": [{"text": "hi"}]},
355                {"role": "model", "parts": [{"text": prose}]},
356            ]
357        });
358        let bytes = serde_json::to_vec(&body).unwrap();
359        let (out, _o, _c) = compress_request_body(body, bytes.len(), None);
360        let parsed: Value = serde_json::from_slice(&out).unwrap();
361
362        assert!(
363            parsed["systemInstruction"]["parts"][0]["text"]
364                .as_str()
365                .unwrap()
366                .len()
367                < prose.len(),
368            "systemInstruction prose must be compressed when enabled"
369        );
370        assert_eq!(
371            parsed["contents"][1]["parts"][0]["text"].as_str().unwrap(),
372            prose,
373            "model (assistant) turns must pass through verbatim (#710)"
374        );
375    }
376
377    #[test]
378    fn gemini_prose_compression_is_deterministic() {
379        let _iso = crate::core::data_dir::isolated_data_dir();
380        crate::core::config::Config::update_global(|c| {
381            c.proxy.role_aggressiveness.system = Some(0.6);
382        })
383        .unwrap();
384        let prose = big_prose();
385        let mk = || {
386            serde_json::json!({
387                "systemInstruction": {"parts": [{"text": prose}]},
388                "contents": [{"role": "user", "parts": [{"text": "hi"}]}]
389            })
390        };
391        let (a, b) = (mk(), mk());
392        let la = serde_json::to_vec(&a).unwrap().len();
393        let lb = serde_json::to_vec(&b).unwrap().len();
394        assert_eq!(
395            compress_request_body(a, la, None).0,
396            compress_request_body(b, lb, None).0,
397            "identical input must yield byte-identical output (#498)"
398        );
399    }
400
401    #[test]
402    fn cache_aware_gemini_prefix_is_byte_stable_across_turns() {
403        // THE cache invariant for the Gemini rail: every `contents` entry before
404        // an already-passed boundary stays byte-identical as the chat grows.
405        let _iso = crate::core::data_dir::isolated_data_dir();
406        let mut prev: Vec<String> = Vec::new();
407        let mut prev_boundary = 0;
408        for pairs in 1..=20 {
409            let contents = gemini_read_turns(pairs);
410            let len = contents.len();
411            let body = serde_json::json!({ "contents": contents });
412            let bytes = serde_json::to_vec(&body).unwrap();
413            let (out, _, _) = compress_request_body(body, bytes.len(), None);
414            let parsed: Value = serde_json::from_slice(&out).unwrap();
415            let items: Vec<String> = parsed["contents"]
416                .as_array()
417                .unwrap()
418                .iter()
419                .map(Value::to_string)
420                .collect();
421            for i in 0..prev_boundary {
422                assert_eq!(
423                    prev[i], items[i],
424                    "Gemini content {i} changed at turn {pairs} — prompt cache prefix broken"
425                );
426            }
427            prev = items;
428            prev_boundary = crate::proxy::history_prune::prune_boundary(
429                crate::core::config::HistoryMode::CacheAware,
430                len,
431            );
432        }
433    }
434
435    #[test]
436    fn effort_control_sets_thinking_config_by_generation() {
437        // #840 end-to-end: the model is taken from the URL path, so the handler
438        // threads it in. 3.x → thinkingLevel; off / unknown model → byte no-op.
439        let _iso = crate::core::data_dir::isolated_data_dir();
440        crate::test_env::remove_var("LEAN_CTX_PROXY_EFFORT");
441        crate::core::config::Config::update_global(|c| {
442            c.proxy.effort = Some("low".into());
443        })
444        .unwrap();
445
446        let body = serde_json::json!({
447            "contents": [{"role": "user", "parts": [{"text": "hi"}]}]
448        });
449        let bytes = serde_json::to_vec(&body).unwrap();
450        let (out, _o, _c) = compress_request_body(body.clone(), bytes.len(), Some("gemini-3-pro"));
451        assert_eq!(
452            serde_json::from_slice::<Value>(&out).unwrap()["generationConfig"]["thinkingConfig"]["thinkingLevel"],
453            "low"
454        );
455
456        // No model (path didn't resolve) → strict no-op, body byte-unchanged.
457        let bytes = serde_json::to_vec(&body).unwrap();
458        let (out, _o, _c) = compress_request_body(body.clone(), bytes.len(), None);
459        assert_eq!(serde_json::from_slice::<Value>(&out).unwrap(), body);
460    }
461}
lean_ctx/proxy/google.rs

lean_ctx/proxy/
google.rs
