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

lean_ctx/proxy/
openai.rs
