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

lean_ctx/proxy/
google.rs
