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(
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
44pub 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 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
71pub(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
114fn 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
142pub(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
175fn 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 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 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 let _iso = crate::core::data_dir::isolated_data_dir();
362 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 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 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 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 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 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}