lean_ctx/proxy/
openai.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.openai_upstream();
20    forward::forward_request(
21        State(state),
22        req,
23        &upstream,
24        "/v1/chat/completions",
25        compress_request_body,
26        "OpenAI",
27        &[],
28    )
29    .await
30}
31
32pub(super) fn compress_request_body(
33    parsed: Value,
34    original_size: usize,
35) -> (Vec<u8>, usize, usize) {
36    let mut doc = parsed;
37    let mut modified = false;
38
39    // Opt-in per-role prose aggressiveness (#710); both default `None` → no-op.
40    let cfg = crate::core::config::Config::load();
41    let system_aggr = cfg.proxy.resolved_role_aggressiveness(ProseRole::System);
42    let user_aggr = cfg.proxy.resolved_role_aggressiveness(ProseRole::User);
43    let live_compress = cfg.proxy.live_compresses();
44    let mode = cfg.proxy.resolved_history_mode();
45    // #895 Track B: output-savings holdout arm, from the pristine body (before any
46    // mutation below) so it matches the arm the response meter records. Control
47    // conversations skip output-shaping but are still metered. Default 0 → Treatment.
48    let arm = super::holdout::assign(
49        &super::holdout::openai_chat_key(&doc),
50        cfg.proxy.output_holdout_fraction(),
51    );
52    // #493: in-band CCR expansion (opt-in). Splice any <lc_expand:HASH> the model
53    // echoed back into the verbatim original from the local tee store. A strict
54    // no-op when no marker is present (byte-identical body → cache-safe). Runs
55    // before the meter-only short-circuit so an explicit expand request is
56    // honored even when the proxy is otherwise byte-passthrough.
57    if cfg.proxy.ccr_inband_enabled() {
58        modified |= super::ccr::splice_inband_in_place(&mut doc);
59    }
60    // #834: cache-safe cross-provider effort control. Default off → no-op. The
61    // value is a constant, so it never perturbs the prompt-cache prefix; it sets
62    // `reasoning_effort` only on reasoning models and never overrides a
63    // client-set value.
64    if arm == super::holdout::Arm::Treatment {
65        if let Some(effort) = cfg.proxy.resolved_effort() {
66            modified |= super::effort::apply_openai_chat(&mut doc, effort);
67        }
68        // #895: cache-safe wire verbosity steer; control arm skips it (measured).
69        if cfg.proxy.verbosity_steer_enabled() {
70            modified |= super::verbosity::apply_openai_chat(&mut doc);
71        }
72    }
73    // Meter-only (#481): nothing rewrites the body, so skip all work and let
74    // forward + usage metering run against the byte-unchanged request. A pending
75    // in-band splice (`modified`) opts out: the body did change this turn.
76    if !live_compress
77        && mode == HistoryMode::Off
78        && system_aggr.is_none()
79        && user_aggr.is_none()
80        && !modified
81    {
82        let out = serde_json::to_vec(&doc).unwrap_or_default();
83        return (out, original_size, original_size);
84    }
85    let mut prose_segments: u64 = 0;
86
87    if let Some(messages) = doc.get_mut("messages").and_then(|m| m.as_array_mut()) {
88        let tool_names = tool_kind::openai_tool_names(messages);
89
90        // OpenAI's automatic prompt caching is prefix-based like Anthropic's,
91        // so history is pruned at the same frozen, cache-aware boundary. `mode`
92        // resolved above.
93        let boundary = super::history_prune::prune_boundary(mode, messages.len());
94        // Mirror the Anthropic guard: never rewrite content behind a client
95        // `cache_control` breakpoint (#448). OpenAI requests carry none, so this
96        // resolves to 0 and pruning is byte-for-byte unchanged — but the code
97        // path stays uniform across providers.
98        let cached = super::history_prune::cached_prefix_len(messages);
99        modified |=
100            super::history_prune::prune_history_range(messages, cached, boundary, &tool_names);
101
102        for msg in messages.iter_mut() {
103            let role = msg.get("role").and_then(|r| r.as_str()).unwrap_or("");
104            if role != "tool" {
105                continue;
106            }
107
108            let name = msg
109                .get("tool_call_id")
110                .and_then(|v| v.as_str())
111                .and_then(|id| tool_names.get(id))
112                .map(String::as_str);
113            let kind = name.map_or(ToolResultKind::Other, tool_kind::classify_tool_name);
114
115            // #481: skip live compression when globally off or tool excluded.
116            if !live_compress || name.is_some_and(|n| cfg.proxy.is_tool_live_compress_excluded(n)) {
117                continue;
118            }
119            if let Some(mut content) = msg
120                .get_mut("content")
121                .and_then(|c| c.as_str().map(String::from))
122                && super::tool_output::compress_text(&mut content, name, kind)
123            {
124                msg["content"] = Value::String(content);
125                modified = true;
126            }
127        }
128
129        // Frozen-region prose: system anchors (deterministic rewrite keeps the
130        // auto-cached prefix byte-stable, so safe at any position outside the
131        // client-cached prefix) and user turns in `[cached, boundary)`. The
132        // `assistant` and `tool` roles are never touched (passthrough). Both
133        // knobs default off, so this is inert unless an operator opts in.
134        if system_aggr.is_some() || user_aggr.is_some() {
135            for (i, msg) in messages.iter_mut().enumerate() {
136                if i < cached {
137                    continue;
138                }
139                let role = msg.get("role").and_then(|r| r.as_str()).unwrap_or("");
140                let aggr = match role {
141                    "system" | "developer" => system_aggr,
142                    "user" if i < boundary => user_aggr,
143                    _ => None,
144                };
145                if let Some(a) = aggr {
146                    prose_segments += u64::from(prose::compress_message_content(msg, a));
147                }
148            }
149        }
150    }
151
152    if prose_segments > 0 {
153        modified = true;
154    }
155    cache_safety::record(prose_segments, true);
156
157    // Ask OpenAI to append a final usage chunk so the proxy can meter real
158    // spend. Not counted as compression (it slightly grows the body), so it
159    // never inflates the savings figure.
160    maybe_inject_usage_reporting(&mut doc);
161
162    let out = serde_json::to_vec(&doc).unwrap_or_default();
163    let compressed_size = if modified { out.len() } else { original_size };
164    (out, original_size, compressed_size)
165}
166
167/// Config-gated wrapper around [`inject_usage_reporting`].
168fn maybe_inject_usage_reporting(doc: &mut Value) {
169    if crate::core::config::Config::load()
170        .proxy
171        .meters_openai_usage()
172    {
173        inject_usage_reporting(doc);
174    }
175}
176
177/// Injects `stream_options.include_usage = true` into a streamed Chat
178/// Completions request so the final SSE chunk carries `usage` (OpenAI omits it
179/// otherwise). No-op for non-streamed requests or when the client already
180/// configured `stream_options.include_usage`.
181fn inject_usage_reporting(doc: &mut Value) {
182    if doc.get("stream").and_then(Value::as_bool) != Some(true) {
183        return;
184    }
185    let Some(obj) = doc.as_object_mut() else {
186        return;
187    };
188    let opts = obj
189        .entry("stream_options")
190        .or_insert_with(|| Value::Object(serde_json::Map::new()));
191    if let Some(opts_obj) = opts.as_object_mut() {
192        opts_obj.entry("include_usage").or_insert(Value::Bool(true));
193    }
194}
195
196#[cfg(test)]
197mod tests {
198    use super::super::compress::compress_tool_result;
199    use super::*;
200
201    #[test]
202    fn read_file_tool_result_protected() {
203        let code = (0..60)
204            .map(|i| format!("    const value{i} = computeValue{i}(ctx, opts);"))
205            .collect::<Vec<_>>()
206            .join("\n");
207        let body = serde_json::json!({
208            "model": "gpt-5",
209            "messages": [
210                {"role": "assistant", "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "read_file"}}]},
211                {"role": "tool", "tool_call_id": "call_1", "content": code}
212            ]
213        });
214        let bytes = serde_json::to_vec(&body).unwrap();
215        let (out, _orig, _comp) = compress_request_body(body, bytes.len());
216        let parsed: Value = serde_json::from_slice(&out).unwrap();
217        assert!(
218            parsed["messages"][1]["content"]
219                .as_str()
220                .unwrap()
221                .contains("value59")
222        );
223    }
224
225    #[test]
226    fn json_envelope_tool_output_is_compressed() {
227        let _iso = crate::core::data_dir::isolated_data_dir();
228        let raw = long_git_status();
229        let expected = compress_tool_result(&raw, Some("Bash"));
230        let envelope = serde_json::to_string(&serde_json::json!({
231            "content": [{"type": "text", "text": raw}],
232            "isError": false,
233        }))
234        .unwrap();
235        let body = serde_json::json!({
236            "model": "gpt-5",
237            "messages": [
238                {"role": "assistant", "tool_calls": [{
239                    "id": "call_1",
240                    "type": "function",
241                    "function": {"name": "Bash"}
242                }]},
243                {"role": "tool", "tool_call_id": "call_1", "content": envelope}
244            ]
245        });
246        let bytes = serde_json::to_vec(&body).unwrap();
247        let (out, orig, comp) = compress_request_body(body, bytes.len());
248
249        assert!(comp < orig, "JSON envelope tool output should shrink");
250        let parsed: Value = serde_json::from_slice(&out).unwrap();
251        let content = parsed["messages"][1]["content"].as_str().unwrap();
252        let envelope: Value = serde_json::from_str(content).unwrap();
253        assert_eq!(envelope["content"][0]["text"].as_str().unwrap(), expected);
254    }
255
256    #[test]
257    fn injects_include_usage_for_streaming() {
258        let mut doc = serde_json::json!({"model": "gpt-5.4", "stream": true, "messages": []});
259        inject_usage_reporting(&mut doc);
260        assert_eq!(doc["stream_options"]["include_usage"], Value::Bool(true));
261    }
262
263    fn long_git_status() -> String {
264        let mut s = String::from(
265            "$ 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",
266        );
267        for i in 0..80 {
268            s.push_str(&format!("\tmodified:   src/module_{i}/file_{i}.rs\n"));
269        }
270        s.push_str("\nno changes added to commit (use \"git add\" and/or \"git commit -a\")\n");
271        s
272    }
273
274    #[test]
275    fn no_injection_for_non_streaming() {
276        let mut doc = serde_json::json!({"model": "gpt-5.4", "messages": []});
277        inject_usage_reporting(&mut doc);
278        assert!(
279            doc.get("stream_options").is_none(),
280            "non-streamed requests get usage in the body, no injection needed"
281        );
282    }
283
284    #[test]
285    fn respects_client_set_include_usage() {
286        let mut doc = serde_json::json!({
287            "model": "gpt-5.4",
288            "stream": true,
289            "stream_options": {"include_usage": false},
290            "messages": []
291        });
292        inject_usage_reporting(&mut doc);
293        assert_eq!(
294            doc["stream_options"]["include_usage"],
295            Value::Bool(false),
296            "an explicit client value must be preserved"
297        );
298    }
299
300    fn big_prose() -> String {
301        let p = "You are a careful, senior software engineer. You always explain your \
302                 reasoning before making changes, you prefer small reviewable diffs, and \
303                 you never introduce mock data or placeholders into production code. ";
304        [p; 6].join("\n")
305    }
306
307    #[test]
308    fn system_message_compressed_and_assistant_untouched() {
309        let _iso = crate::core::data_dir::isolated_data_dir();
310        crate::core::config::Config::update_global(|c| {
311            c.proxy.role_aggressiveness.system = Some(0.6);
312        })
313        .unwrap();
314
315        let prose = big_prose();
316        let body = serde_json::json!({
317            "model": "gpt-5",
318            "messages": [
319                {"role": "system", "content": prose},
320                {"role": "user", "content": "hi"},
321                {"role": "assistant", "content": prose},
322            ]
323        });
324        let bytes = serde_json::to_vec(&body).unwrap();
325        let (out, _o, _c) = compress_request_body(body, bytes.len());
326        let parsed: Value = serde_json::from_slice(&out).unwrap();
327
328        assert!(
329            parsed["messages"][0]["content"].as_str().unwrap().len() < prose.len(),
330            "system message prose must be compressed when enabled"
331        );
332        assert_eq!(
333            parsed["messages"][2]["content"].as_str().unwrap(),
334            prose,
335            "assistant turns must pass through verbatim (#710)"
336        );
337    }
338
339    #[test]
340    fn openai_prose_compression_is_deterministic() {
341        let _iso = crate::core::data_dir::isolated_data_dir();
342        crate::core::config::Config::update_global(|c| {
343            c.proxy.role_aggressiveness.system = Some(0.6);
344        })
345        .unwrap();
346        let prose = big_prose();
347        let mk = || {
348            serde_json::json!({
349                "model": "gpt-5",
350                "messages": [{"role": "system", "content": prose}, {"role": "user", "content": "hi"}]
351            })
352        };
353        let (a, b) = (mk(), mk());
354        let la = serde_json::to_vec(&a).unwrap().len();
355        let lb = serde_json::to_vec(&b).unwrap().len();
356        assert_eq!(
357            compress_request_body(a, la).0,
358            compress_request_body(b, lb).0,
359            "identical input must yield byte-identical output (#498)"
360        );
361    }
362
363    #[test]
364    fn effort_control_sets_reasoning_effort_and_off_is_noop() {
365        // #834 end-to-end through the Chat Completions request path.
366        let _iso = crate::core::data_dir::isolated_data_dir();
367        crate::test_env::remove_var("LEAN_CTX_PROXY_EFFORT");
368        let body = serde_json::json!({
369            "model": "gpt-5.5", "messages": [{"role": "user", "content": "hi"}]
370        });
371        let bytes = serde_json::to_vec(&body).unwrap();
372
373        // Off by default: a cache-safe no-op (size unchanged, no param added).
374        let (off, o, c) = compress_request_body(body.clone(), bytes.len());
375        assert_eq!(c, o, "effort off must be a passthrough");
376        assert!(
377            serde_json::from_slice::<Value>(&off)
378                .unwrap()
379                .get("reasoning_effort")
380                .is_none()
381        );
382
383        // Enabled: reasoning_effort is filled on the reasoning model.
384        crate::core::config::Config::update_global(|cfg| {
385            cfg.proxy.effort = Some("low".into());
386        })
387        .unwrap();
388        let (on, _o, _c) = compress_request_body(body, bytes.len());
389        assert_eq!(
390            serde_json::from_slice::<Value>(&on).unwrap()["reasoning_effort"],
391            "low"
392        );
393    }
394}
lean_ctx/proxy/openai.rs

lean_ctx/proxy/
openai.rs
