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

lean_ctx/proxy/
google.rs
