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};
13
14pub async fn handler(
15    State(state): State<ProxyState>,
16    req: Request<Body>,
17) -> Result<Response, StatusCode> {
18    let upstream = state.openai_upstream();
19    forward::forward_request(
20        State(state),
21        req,
22        &upstream,
23        "/v1/chat/completions",
24        compress_request_body,
25        "OpenAI",
26        &[],
27    )
28    .await
29}
30
31fn compress_request_body(parsed: Value, original_size: usize) -> (Vec<u8>, usize, usize) {
32    let mut doc = parsed;
33    let mut modified = false;
34
35    if let Some(messages) = doc.get_mut("messages").and_then(|m| m.as_array_mut()) {
36        let tool_names = tool_kind::openai_tool_names(messages);
37
38        // OpenAI's automatic prompt caching is prefix-based like Anthropic's,
39        // so history is pruned at the same frozen, cache-aware boundary.
40        let mode = crate::core::config::Config::load()
41            .proxy
42            .resolved_history_mode();
43        let boundary = super::history_prune::prune_boundary(mode, messages.len());
44        // Mirror the Anthropic guard: never rewrite content behind a client
45        // `cache_control` breakpoint (#448). OpenAI requests carry none, so this
46        // resolves to 0 and pruning is byte-for-byte unchanged — but the code
47        // path stays uniform across providers.
48        let cached = super::history_prune::cached_prefix_len(messages);
49        modified |=
50            super::history_prune::prune_history_range(messages, cached, boundary, &tool_names);
51
52        for msg in messages.iter_mut() {
53            let role = msg.get("role").and_then(|r| r.as_str()).unwrap_or("");
54            if role != "tool" {
55                continue;
56            }
57
58            let name = msg
59                .get("tool_call_id")
60                .and_then(|v| v.as_str())
61                .and_then(|id| tool_names.get(id))
62                .map(String::as_str);
63            let kind = name.map_or(ToolResultKind::Other, tool_kind::classify_tool_name);
64
65            if let Some(content) = msg
66                .get_mut("content")
67                .and_then(|c| c.as_str().map(String::from))
68            {
69                if should_protect(kind, &content) {
70                    continue;
71                }
72                let compressed = compress_tool_result(&content, name);
73                if compressed.len() < content.len() {
74                    msg["content"] = Value::String(compressed);
75                    modified = true;
76                }
77            }
78        }
79    }
80
81    // Ask OpenAI to append a final usage chunk so the proxy can meter real
82    // spend. Not counted as compression (it slightly grows the body), so it
83    // never inflates the savings figure.
84    maybe_inject_usage_reporting(&mut doc);
85
86    let out = serde_json::to_vec(&doc).unwrap_or_default();
87    let compressed_size = if modified { out.len() } else { original_size };
88    (out, original_size, compressed_size)
89}
90
91/// Config-gated wrapper around [`inject_usage_reporting`].
92fn maybe_inject_usage_reporting(doc: &mut Value) {
93    if crate::core::config::Config::load()
94        .proxy
95        .meters_openai_usage()
96    {
97        inject_usage_reporting(doc);
98    }
99}
100
101/// Injects `stream_options.include_usage = true` into a streamed Chat
102/// Completions request so the final SSE chunk carries `usage` (OpenAI omits it
103/// otherwise). No-op for non-streamed requests or when the client already
104/// configured `stream_options.include_usage`.
105fn inject_usage_reporting(doc: &mut Value) {
106    if doc.get("stream").and_then(Value::as_bool) != Some(true) {
107        return;
108    }
109    let Some(obj) = doc.as_object_mut() else {
110        return;
111    };
112    let opts = obj
113        .entry("stream_options")
114        .or_insert_with(|| Value::Object(serde_json::Map::new()));
115    if let Some(opts_obj) = opts.as_object_mut() {
116        opts_obj.entry("include_usage").or_insert(Value::Bool(true));
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123
124    #[test]
125    fn read_file_tool_result_protected() {
126        let code = (0..60)
127            .map(|i| format!("    const value{i} = computeValue{i}(ctx, opts);"))
128            .collect::<Vec<_>>()
129            .join("\n");
130        let body = serde_json::json!({
131            "model": "gpt-5",
132            "messages": [
133                {"role": "assistant", "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "read_file"}}]},
134                {"role": "tool", "tool_call_id": "call_1", "content": code}
135            ]
136        });
137        let bytes = serde_json::to_vec(&body).unwrap();
138        let (out, _orig, _comp) = compress_request_body(body, bytes.len());
139        let parsed: Value = serde_json::from_slice(&out).unwrap();
140        assert!(
141            parsed["messages"][1]["content"]
142                .as_str()
143                .unwrap()
144                .contains("value59")
145        );
146    }
147
148    #[test]
149    fn injects_include_usage_for_streaming() {
150        let mut doc = serde_json::json!({"model": "gpt-5.4", "stream": true, "messages": []});
151        inject_usage_reporting(&mut doc);
152        assert_eq!(doc["stream_options"]["include_usage"], Value::Bool(true));
153    }
154
155    #[test]
156    fn no_injection_for_non_streaming() {
157        let mut doc = serde_json::json!({"model": "gpt-5.4", "messages": []});
158        inject_usage_reporting(&mut doc);
159        assert!(
160            doc.get("stream_options").is_none(),
161            "non-streamed requests get usage in the body, no injection needed"
162        );
163    }
164
165    #[test]
166    fn respects_client_set_include_usage() {
167        let mut doc = serde_json::json!({
168            "model": "gpt-5.4",
169            "stream": true,
170            "stream_options": {"include_usage": false},
171            "messages": []
172        });
173        inject_usage_reporting(&mut doc);
174        assert_eq!(
175            doc["stream_options"]["include_usage"],
176            Value::Bool(false),
177            "an explicit client value must be preserved"
178        );
179    }
180}
lean_ctx/proxy/openai.rs

lean_ctx/proxy/
openai.rs
