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::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 mut prose_segments: u64 = 0;
42
43    if let Some(messages) = doc.get_mut("messages").and_then(|m| m.as_array_mut()) {
44        let tool_names = tool_kind::openai_tool_names(messages);
45
46        // OpenAI's automatic prompt caching is prefix-based like Anthropic's,
47        // so history is pruned at the same frozen, cache-aware boundary.
48        let mode = cfg.proxy.resolved_history_mode();
49        let boundary = super::history_prune::prune_boundary(mode, messages.len());
50        // Mirror the Anthropic guard: never rewrite content behind a client
51        // `cache_control` breakpoint (#448). OpenAI requests carry none, so this
52        // resolves to 0 and pruning is byte-for-byte unchanged — but the code
53        // path stays uniform across providers.
54        let cached = super::history_prune::cached_prefix_len(messages);
55        modified |=
56            super::history_prune::prune_history_range(messages, cached, boundary, &tool_names);
57
58        for msg in messages.iter_mut() {
59            let role = msg.get("role").and_then(|r| r.as_str()).unwrap_or("");
60            if role != "tool" {
61                continue;
62            }
63
64            let name = msg
65                .get("tool_call_id")
66                .and_then(|v| v.as_str())
67                .and_then(|id| tool_names.get(id))
68                .map(String::as_str);
69            let kind = name.map_or(ToolResultKind::Other, tool_kind::classify_tool_name);
70
71            if let Some(content) = msg
72                .get_mut("content")
73                .and_then(|c| c.as_str().map(String::from))
74            {
75                if should_protect(kind, &content) {
76                    continue;
77                }
78                let compressed = compress_tool_result(&content, name);
79                if compressed.len() < content.len() {
80                    msg["content"] = Value::String(compressed);
81                    modified = true;
82                }
83            }
84        }
85
86        // Frozen-region prose: system anchors (deterministic rewrite keeps the
87        // auto-cached prefix byte-stable, so safe at any position outside the
88        // client-cached prefix) and user turns in `[cached, boundary)`. The
89        // `assistant` and `tool` roles are never touched (passthrough). Both
90        // knobs default off, so this is inert unless an operator opts in.
91        if system_aggr.is_some() || user_aggr.is_some() {
92            for (i, msg) in messages.iter_mut().enumerate() {
93                if i < cached {
94                    continue;
95                }
96                let role = msg.get("role").and_then(|r| r.as_str()).unwrap_or("");
97                let aggr = match role {
98                    "system" | "developer" => system_aggr,
99                    "user" if i < boundary => user_aggr,
100                    _ => None,
101                };
102                if let Some(a) = aggr {
103                    prose_segments += u64::from(prose::compress_message_content(msg, a));
104                }
105            }
106        }
107    }
108
109    if prose_segments > 0 {
110        modified = true;
111    }
112    cache_safety::record(prose_segments, true);
113
114    // Ask OpenAI to append a final usage chunk so the proxy can meter real
115    // spend. Not counted as compression (it slightly grows the body), so it
116    // never inflates the savings figure.
117    maybe_inject_usage_reporting(&mut doc);
118
119    let out = serde_json::to_vec(&doc).unwrap_or_default();
120    let compressed_size = if modified { out.len() } else { original_size };
121    (out, original_size, compressed_size)
122}
123
124/// Config-gated wrapper around [`inject_usage_reporting`].
125fn maybe_inject_usage_reporting(doc: &mut Value) {
126    if crate::core::config::Config::load()
127        .proxy
128        .meters_openai_usage()
129    {
130        inject_usage_reporting(doc);
131    }
132}
133
134/// Injects `stream_options.include_usage = true` into a streamed Chat
135/// Completions request so the final SSE chunk carries `usage` (OpenAI omits it
136/// otherwise). No-op for non-streamed requests or when the client already
137/// configured `stream_options.include_usage`.
138fn inject_usage_reporting(doc: &mut Value) {
139    if doc.get("stream").and_then(Value::as_bool) != Some(true) {
140        return;
141    }
142    let Some(obj) = doc.as_object_mut() else {
143        return;
144    };
145    let opts = obj
146        .entry("stream_options")
147        .or_insert_with(|| Value::Object(serde_json::Map::new()));
148    if let Some(opts_obj) = opts.as_object_mut() {
149        opts_obj.entry("include_usage").or_insert(Value::Bool(true));
150    }
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156
157    #[test]
158    fn read_file_tool_result_protected() {
159        let code = (0..60)
160            .map(|i| format!("    const value{i} = computeValue{i}(ctx, opts);"))
161            .collect::<Vec<_>>()
162            .join("\n");
163        let body = serde_json::json!({
164            "model": "gpt-5",
165            "messages": [
166                {"role": "assistant", "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "read_file"}}]},
167                {"role": "tool", "tool_call_id": "call_1", "content": code}
168            ]
169        });
170        let bytes = serde_json::to_vec(&body).unwrap();
171        let (out, _orig, _comp) = compress_request_body(body, bytes.len());
172        let parsed: Value = serde_json::from_slice(&out).unwrap();
173        assert!(
174            parsed["messages"][1]["content"]
175                .as_str()
176                .unwrap()
177                .contains("value59")
178        );
179    }
180
181    #[test]
182    fn injects_include_usage_for_streaming() {
183        let mut doc = serde_json::json!({"model": "gpt-5.4", "stream": true, "messages": []});
184        inject_usage_reporting(&mut doc);
185        assert_eq!(doc["stream_options"]["include_usage"], Value::Bool(true));
186    }
187
188    #[test]
189    fn no_injection_for_non_streaming() {
190        let mut doc = serde_json::json!({"model": "gpt-5.4", "messages": []});
191        inject_usage_reporting(&mut doc);
192        assert!(
193            doc.get("stream_options").is_none(),
194            "non-streamed requests get usage in the body, no injection needed"
195        );
196    }
197
198    #[test]
199    fn respects_client_set_include_usage() {
200        let mut doc = serde_json::json!({
201            "model": "gpt-5.4",
202            "stream": true,
203            "stream_options": {"include_usage": false},
204            "messages": []
205        });
206        inject_usage_reporting(&mut doc);
207        assert_eq!(
208            doc["stream_options"]["include_usage"],
209            Value::Bool(false),
210            "an explicit client value must be preserved"
211        );
212    }
213
214    fn big_prose() -> String {
215        let p = "You are a careful, senior software engineer. You always explain your \
216                 reasoning before making changes, you prefer small reviewable diffs, and \
217                 you never introduce mock data or placeholders into production code. ";
218        [p; 6].join("\n")
219    }
220
221    #[test]
222    fn system_message_compressed_and_assistant_untouched() {
223        let _iso = crate::core::data_dir::isolated_data_dir();
224        crate::core::config::Config::update_global(|c| {
225            c.proxy.role_aggressiveness.system = Some(0.6);
226        })
227        .unwrap();
228
229        let prose = big_prose();
230        let body = serde_json::json!({
231            "model": "gpt-5",
232            "messages": [
233                {"role": "system", "content": prose},
234                {"role": "user", "content": "hi"},
235                {"role": "assistant", "content": prose},
236            ]
237        });
238        let bytes = serde_json::to_vec(&body).unwrap();
239        let (out, _o, _c) = compress_request_body(body, bytes.len());
240        let parsed: Value = serde_json::from_slice(&out).unwrap();
241
242        assert!(
243            parsed["messages"][0]["content"].as_str().unwrap().len() < prose.len(),
244            "system message prose must be compressed when enabled"
245        );
246        assert_eq!(
247            parsed["messages"][2]["content"].as_str().unwrap(),
248            prose,
249            "assistant turns must pass through verbatim (#710)"
250        );
251    }
252
253    #[test]
254    fn openai_prose_compression_is_deterministic() {
255        let _iso = crate::core::data_dir::isolated_data_dir();
256        crate::core::config::Config::update_global(|c| {
257            c.proxy.role_aggressiveness.system = Some(0.6);
258        })
259        .unwrap();
260        let prose = big_prose();
261        let mk = || {
262            serde_json::json!({
263                "model": "gpt-5",
264                "messages": [{"role": "system", "content": prose}, {"role": "user", "content": "hi"}]
265            })
266        };
267        let (a, b) = (mk(), mk());
268        let la = serde_json::to_vec(&a).unwrap().len();
269        let lb = serde_json::to_vec(&b).unwrap().len();
270        assert_eq!(
271            compress_request_body(a, la).0,
272            compress_request_body(b, lb).0,
273            "identical input must yield byte-identical output (#498)"
274        );
275    }
276}
lean_ctx/proxy/openai.rs

lean_ctx/proxy/
openai.rs
