Skip to main content

lean_ctx/proxy/
openai_responses.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
14/// Proxy handler for OpenAI's Responses API (`POST /v1/responses`).
15///
16/// The Responses API superseded Chat Completions for clients such as opencode
17/// and the OpenAI Agents SDK. Its conversation turns live in `input` rather than
18/// `messages`, so the Chat Completions handler never saw — and never compressed —
19/// them. This handler reuses the same upstream, auth and streaming path but
20/// understands the Responses-API request shape.
21///
22/// Retrieve / cancel / delete / input_items sub-paths
23/// (`/v1/responses/{id}/...`) are routed here as well and pass through untouched:
24/// they carry no `input` array, so `compress_request_body` is a no-op for them.
25///
26/// Handles `POST /v1/responses` (and the bare `/responses`) over HTTP/SSE.
27pub async fn handler(
28    State(state): State<ProxyState>,
29    req: Request<Body>,
30) -> Result<Response, StatusCode> {
31    let upstream = state.openai_upstream();
32    forward::forward_request(
33        State(state),
34        req,
35        &upstream,
36        "/v1/responses",
37        compress_request_body,
38        "OpenAI",
39        &[],
40    )
41    .await
42}
43
44/// Handles the WebSocket Responses transport on `GET /v1/responses`.
45///
46/// Codex (and the OpenAI SDK) default to `ws://…/responses` with one
47/// `response.create` event per turn. Bridging the upgrade here lets the proxy be
48/// a drop-in for Codex without forcing `supports_websockets = false` (#440); the
49/// actual WS↔HTTP/SSE bridging lives in `openai_responses_ws`.
50pub async fn ws_handler(
51    State(state): State<ProxyState>,
52    headers: axum::http::HeaderMap,
53    ws: axum::extract::ws::WebSocketUpgrade,
54) -> Response {
55    super::openai_responses_ws::upgrade(state, ws, &headers)
56}
57
58fn compress_request_body(parsed: Value, original_size: usize) -> (Vec<u8>, usize, usize) {
59    let mut doc = parsed;
60    // Two-stage, like the Chat Completions path: (1) cache-aware prune of the
61    // frozen OLD region — old file reads collapse to re-read stubs, old logs
62    // head/tail summarize — then (2) compress whatever recent outputs remain.
63    // Stage 1 runs first so a stubbed old output isn't needlessly re-compressed.
64    let mut modified = prune_responses_input(&mut doc);
65    modified |= compress_responses_input(&mut doc);
66    let out = serde_json::to_vec(&doc).unwrap_or_default();
67    let compressed_size = if modified { out.len() } else { original_size };
68    (out, original_size, compressed_size)
69}
70
71/// Cache-aware history pruning for the Responses API.
72///
73/// Unlike the Chat Completions path we never *remove* an item: the Responses API
74/// rejects a `function_call` whose matching `function_call_output` is absent (and
75/// reasoning items must keep their originating call). Instead we rewrite the
76/// `output` text of every `function_call_output` in the frozen OLD region
77/// (`input[..boundary]`) — pairing and ordering are untouched, so there is no
78/// risk of a 400.
79///
80/// The boundary is the same monotone staircase as every other rail
81/// ([`history_prune::prune_boundary`]), so the request prefix stays byte-stable
82/// for up to a full stride and OpenAI's automatic prompt cache keeps hitting.
83///
84/// Shared with the WebSocket bridge (#440) so Codex/WS turns prune identically.
85pub(super) fn prune_responses_input(doc: &mut Value) -> bool {
86    let mode = crate::core::config::Config::load()
87        .proxy
88        .resolved_history_mode();
89    let Some(input) = doc.get_mut("input").and_then(|i| i.as_array_mut()) else {
90        return false;
91    };
92    let boundary = super::history_prune::prune_boundary(mode, input.len());
93    if boundary == 0 {
94        return false;
95    }
96    let tool_names = tool_kind::responses_tool_names(input);
97    let mut modified = false;
98    for item in input.iter_mut().take(boundary) {
99        if item.get("type").and_then(|t| t.as_str()) != Some("function_call_output") {
100            continue;
101        }
102        let kind = item
103            .get("call_id")
104            .and_then(|v| v.as_str())
105            .and_then(|id| tool_names.get(id))
106            .map_or(ToolResultKind::Other, |n| tool_kind::classify_tool_name(n));
107        if let Some(output) = item.get_mut("output") {
108            modified |= prune_output_field(output, kind);
109        }
110    }
111    modified
112}
113
114/// Apply [`history_prune::prune_output_text`] to a `function_call_output.output`,
115/// handling both the JSON-string and array-of-content-parts shapes — the
116/// pruning analogue of [`compress_output_field`].
117fn prune_output_field(output: &mut Value, kind: ToolResultKind) -> bool {
118    match output {
119        Value::String(s) => match super::history_prune::prune_output_text(s, kind) {
120            Some(pruned) => {
121                *s = pruned;
122                true
123            }
124            None => false,
125        },
126        Value::Array(parts) => {
127            let mut changed = false;
128            for part in parts.iter_mut() {
129                if let Some(Value::String(text)) = part.get_mut("text")
130                    && let Some(pruned) = super::history_prune::prune_output_text(text, kind)
131                {
132                    *text = pruned;
133                    changed = true;
134                }
135            }
136            changed
137        }
138        _ => false,
139    }
140}
141
142/// Compresses the `function_call_output.output` entries of a Responses-API body
143/// in place, returning whether anything changed. Shared by the HTTP handler and
144/// the WebSocket bridge (#440) so both paths get identical, safe savings.
145///
146/// The only token sink we shrink is each `function_call_output.output` — the
147/// Responses-API analogue of a Chat Completions `role:"tool"` message. We never
148/// remove or reorder `input` items: the Responses API rejects a `function_call`
149/// whose matching `function_call_output` is absent (and reasoning items must keep
150/// their originating call), so all token reclamation happens *in place* on the
151/// output text. Cache-aware pruning of the frozen OLD region lives in
152/// [`prune_responses_input`]; this pass compresses whatever recent outputs remain.
153pub(super) fn compress_responses_input(doc: &mut Value) -> bool {
154    let mut modified = false;
155    if let Some(input) = doc.get_mut("input").and_then(|i| i.as_array_mut()) {
156        let tool_names = tool_kind::responses_tool_names(input);
157        for item in input.iter_mut() {
158            if item.get("type").and_then(|t| t.as_str()) != Some("function_call_output") {
159                continue;
160            }
161            let name = item
162                .get("call_id")
163                .and_then(|v| v.as_str())
164                .and_then(|id| tool_names.get(id))
165                .map(String::as_str);
166            let kind = name.map_or(ToolResultKind::Other, tool_kind::classify_tool_name);
167            if let Some(output) = item.get_mut("output") {
168                modified |= compress_output_field(output, name, kind);
169            }
170        }
171    }
172    modified
173}
174
175/// Compress a `function_call_output.output`. OpenAI sends this as a JSON string,
176/// but the API also accepts an array of content parts (`input_text` blocks) for
177/// tools returning richer data, so both shapes are handled.
178///
179/// A protected file/source read (resolved from the matching `function_call`
180/// name) is left intact so a mid-refactor model never loses the body it edits.
181fn compress_output_field(
182    output: &mut Value,
183    tool_name: Option<&str>,
184    kind: ToolResultKind,
185) -> bool {
186    match output {
187        Value::String(s) => {
188            if should_protect(kind, s) {
189                return false;
190            }
191            let compressed = compress_tool_result(s, tool_name);
192            if compressed.len() < s.len() {
193                *s = compressed;
194                return true;
195            }
196            false
197        }
198        Value::Array(parts) => {
199            let mut changed = false;
200            for part in parts.iter_mut() {
201                if let Some(Value::String(text)) = part.get_mut("text") {
202                    if should_protect(kind, text) {
203                        continue;
204                    }
205                    let compressed = compress_tool_result(text, tool_name);
206                    if compressed.len() < text.len() {
207                        *text = compressed;
208                        changed = true;
209                    }
210                }
211            }
212            changed
213        }
214        _ => false,
215    }
216}
217
218#[cfg(test)]
219mod tests {
220    use super::*;
221
222    /// A long `git status` is a known-compressible fixture: `has_structural_output`
223    /// is false for it, so it flows through the git-status pattern compressor.
224    fn long_git_status() -> String {
225        let mut s = String::from(
226            "$ 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",
227        );
228        for i in 0..80 {
229            s.push_str(&format!("\tmodified:   src/module_{i}/file_{i}.rs\n"));
230        }
231        s.push_str("\nno changes added to commit (use \"git add\" and/or \"git commit -a\")\n");
232        s
233    }
234
235    #[test]
236    fn string_output_mirrors_engine_and_shrinks() {
237        let raw = long_git_status();
238        let expected = compress_tool_result(&raw, None);
239        assert!(
240            expected.len() < raw.len(),
241            "fixture must be compressible by the shared engine"
242        );
243
244        let body = serde_json::json!({
245            "model": "gpt-5",
246            "input": [
247                {"type": "function_call_output", "call_id": "call_1", "output": raw}
248            ]
249        });
250        let bytes = serde_json::to_vec(&body).unwrap();
251        let (out, orig, comp) = compress_request_body(body, bytes.len());
252
253        assert!(comp < orig, "compressed body must be smaller");
254        let parsed: Value = serde_json::from_slice(&out).unwrap();
255        assert_eq!(
256            parsed["input"][0]["output"].as_str().unwrap(),
257            expected,
258            "output must be exactly what the shared compressor produces"
259        );
260    }
261
262    #[test]
263    fn array_output_text_is_compressed() {
264        let raw = long_git_status();
265        let expected = compress_tool_result(&raw, None);
266
267        let body = serde_json::json!({
268            "input": [
269                {
270                    "type": "function_call_output",
271                    "call_id": "call_1",
272                    "output": [{"type": "input_text", "text": raw}]
273                }
274            ]
275        });
276        let bytes = serde_json::to_vec(&body).unwrap();
277        let (out, orig, comp) = compress_request_body(body, bytes.len());
278
279        assert!(comp < orig);
280        let parsed: Value = serde_json::from_slice(&out).unwrap();
281        assert_eq!(
282            parsed["input"][0]["output"][0]["text"].as_str().unwrap(),
283            expected
284        );
285    }
286
287    #[test]
288    fn non_tool_output_items_are_untouched() {
289        let body = serde_json::json!({
290            "input": [
291                {"type": "message", "role": "user", "content": long_git_status()},
292                {"type": "function_call", "call_id": "c", "name": "x", "arguments": "{}"}
293            ]
294        });
295        let bytes = serde_json::to_vec(&body).unwrap();
296        let (out, orig, comp) = compress_request_body(body.clone(), bytes.len());
297
298        assert_eq!(comp, orig, "no function_call_output → passthrough");
299        let reparsed: Value = serde_json::from_slice(&out).unwrap();
300        assert_eq!(reparsed, body);
301    }
302
303    #[test]
304    fn plain_string_input_passthrough() {
305        let body = serde_json::json!({"model": "gpt-5", "input": "hello world"});
306        let bytes = serde_json::to_vec(&body).unwrap();
307        let (out, orig, comp) = compress_request_body(body.clone(), bytes.len());
308        assert_eq!(comp, orig);
309        let reparsed: Value = serde_json::from_slice(&out).unwrap();
310        assert_eq!(reparsed, body);
311    }
312
313    #[test]
314    fn no_input_field_passthrough() {
315        let body = serde_json::json!({"model": "gpt-5", "previous_response_id": "resp_abc"});
316        let bytes = serde_json::to_vec(&body).unwrap();
317        let (out, orig, comp) = compress_request_body(body.clone(), bytes.len());
318        assert_eq!(comp, orig);
319        let reparsed: Value = serde_json::from_slice(&out).unwrap();
320        assert_eq!(reparsed, body);
321    }
322
323    #[test]
324    fn short_output_unchanged() {
325        let body = serde_json::json!({
326            "input": [
327                {"type": "function_call_output", "call_id": "c", "output": "ok"}
328            ]
329        });
330        let bytes = serde_json::to_vec(&body).unwrap();
331        let (out, orig, comp) = compress_request_body(body.clone(), bytes.len());
332        assert_eq!(comp, orig);
333        let reparsed: Value = serde_json::from_slice(&out).unwrap();
334        assert_eq!(reparsed, body);
335    }
336
337    /// `pairs` Responses turns: each is a `function_call` + its matching
338    /// `function_call_output` carrying a long file read.
339    fn responses_read_turns(pairs: usize) -> Vec<Value> {
340        let code = (0..40)
341            .map(|i| format!("    let v{i} = compute_{i}(ctx, opts);"))
342            .collect::<Vec<_>>()
343            .join("\n");
344        let mut input = Vec::new();
345        for t in 0..pairs {
346            input.push(serde_json::json!({
347                "type": "function_call", "call_id": format!("c{t}"),
348                "name": "read_file", "arguments": "{}"
349            }));
350            input.push(serde_json::json!({
351                "type": "function_call_output", "call_id": format!("c{t}"),
352                "output": format!("{code}\n// turn {t}")
353            }));
354        }
355        input
356    }
357
358    #[test]
359    fn cache_aware_prune_stubs_old_reads_keeps_recent_and_pairing() {
360        // Default (isolated) config = cache-aware history mode.
361        let _iso = crate::core::data_dir::isolated_data_dir();
362        // 14 pairs = 28 items → staircase boundary 16.
363        let body = serde_json::json!({"model": "gpt-5", "input": responses_read_turns(14)});
364        let item_count = body["input"].as_array().unwrap().len();
365        let bytes = serde_json::to_vec(&body).unwrap();
366        let (out, orig, comp) = compress_request_body(body, bytes.len());
367        assert!(comp < orig, "old reads must be pruned for savings");
368
369        let parsed: Value = serde_json::from_slice(&out).unwrap();
370        let input = parsed["input"].as_array().unwrap();
371        // Pairing + ordering preserved: not a single item dropped or moved.
372        assert_eq!(input.len(), item_count, "no items may be removed (pairing)");
373        for (i, item) in input.iter().enumerate() {
374            let expect = if i.is_multiple_of(2) {
375                "function_call"
376            } else {
377                "function_call_output"
378            };
379            assert_eq!(item["type"], expect, "item {i} type/order changed");
380        }
381        // An OLD file read (output index 1, before boundary 16) is stubbed.
382        let old = input[1]["output"].as_str().unwrap();
383        assert!(
384            old.contains("Re-read the file"),
385            "old read should be stubbed, got: {old}"
386        );
387        // A RECENT file read (output index 27, after the boundary) keeps its body.
388        let recent = input[27]["output"].as_str().unwrap();
389        assert!(
390            recent.contains("v39"),
391            "recent read must be protected, got: {recent}"
392        );
393    }
394
395    #[test]
396    fn responses_compression_is_deterministic() {
397        // #498: the same request must compress to byte-identical output so the
398        // provider's prompt cache (and our regression diffs) stay stable.
399        let _iso = crate::core::data_dir::isolated_data_dir();
400        let mk = || serde_json::json!({"model": "gpt-5", "input": responses_read_turns(14)});
401        let (a, b) = (mk(), mk());
402        let (la, lb) = (
403            serde_json::to_vec(&a).unwrap().len(),
404            serde_json::to_vec(&b).unwrap().len(),
405        );
406        let (out_a, _, _) = compress_request_body(a, la);
407        let (out_b, _, _) = compress_request_body(b, lb);
408        assert_eq!(out_a, out_b, "identical input must yield identical bytes");
409    }
410
411    #[test]
412    fn cache_aware_responses_prefix_is_byte_stable_across_turns() {
413        // THE cache invariant for the Responses rail: as `input` grows turn by
414        // turn, every item before an already-passed boundary must stay
415        // byte-identical, or OpenAI's automatic prompt cache stops hitting.
416        let _iso = crate::core::data_dir::isolated_data_dir();
417        let mut prev: Vec<String> = Vec::new();
418        let mut prev_boundary = 0;
419        for pairs in 1..=20 {
420            let input = responses_read_turns(pairs);
421            let len = input.len();
422            let body = serde_json::json!({"model": "gpt-5", "input": input});
423            let bytes = serde_json::to_vec(&body).unwrap();
424            let (out, _, _) = compress_request_body(body, bytes.len());
425            let parsed: Value = serde_json::from_slice(&out).unwrap();
426            let items: Vec<String> = parsed["input"]
427                .as_array()
428                .unwrap()
429                .iter()
430                .map(Value::to_string)
431                .collect();
432            for i in 0..prev_boundary {
433                assert_eq!(
434                    prev[i], items[i],
435                    "Responses item {i} changed at turn {pairs} — prompt cache prefix broken"
436                );
437            }
438            prev = items;
439            prev_boundary = crate::proxy::history_prune::prune_boundary(
440                crate::core::config::HistoryMode::CacheAware,
441                len,
442            );
443        }
444    }
445}