Skip to main content

lean_ctx/proxy/
anthropic.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.anthropic_upstream();
21    forward::forward_request(
22        State(state),
23        req,
24        &upstream,
25        "/v1/messages",
26        compress_request_body,
27        "Anthropic",
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 to `None`, in
38    // which case nothing below fires and the body is byte-for-byte unchanged.
39    let cfg = crate::core::config::Config::load();
40    let system_aggr = cfg.proxy.resolved_role_aggressiveness(ProseRole::System);
41    let user_aggr = cfg.proxy.resolved_role_aggressiveness(ProseRole::User);
42    let mut prose_segments: u64 = 0;
43
44    // Length of the client's provider-cached message prefix. Needed both for
45    // cache-safe pruning below and to gate top-level system prose: if any
46    // message is client-cached, `system` (which precedes every message) is part
47    // of that cached prefix and must not be rewritten.
48    let cached = doc
49        .get("messages")
50        .and_then(|m| m.as_array())
51        .map_or(0, |m| super::history_prune::cached_prefix_len(m));
52
53    // System prose: only when nothing is client-cached and the `system` field
54    // carries no `cache_control` of its own — otherwise it anchors the cache.
55    if let Some(a) = system_aggr
56        && cached == 0
57        && let Some(system) = doc.get_mut("system")
58        && !prose::value_has_cache_control(system)
59    {
60        let n = prose::compress_system_value(system, a);
61        if n > 0 {
62            prose_segments += u64::from(n);
63            modified = true;
64        }
65    }
66
67    if let Some(messages) = doc.get_mut("messages").and_then(|m| m.as_array_mut()) {
68        // Resolve tool-call id → tool name so file/source reads can be protected
69        // from lossy compression that would force the model to re-read mid-task.
70        let tool_names = tool_kind::anthropic_tool_names(messages);
71
72        // Prune at a frozen, cache-aware boundary by default: Anthropic's
73        // prompt cache matches exact prefixes, so the boundary must not move
74        // every turn (see `history_prune::prune_boundary`).
75        let mode = cfg.proxy.resolved_history_mode();
76        let boundary = super::history_prune::prune_boundary(mode, messages.len());
77        // Never rewrite content the client has marked with `cache_control`:
78        // pruning inside the already-cached prefix invalidates Anthropic's
79        // prompt cache from the first changed message (#448). Pruning therefore
80        // starts after the last breakpoint; with no breakpoint this is 0, i.e.
81        // the previous behaviour.
82        modified |=
83            super::history_prune::prune_history_range(messages, cached, boundary, &tool_names);
84
85        for msg in messages.iter_mut() {
86            let role = msg.get("role").and_then(|r| r.as_str()).unwrap_or("");
87            if role != "user" {
88                continue;
89            }
90
91            if let Some(content) = msg.get_mut("content").and_then(|c| c.as_array_mut()) {
92                for block in content.iter_mut() {
93                    if block.get("type").and_then(|t| t.as_str()) != Some("tool_result") {
94                        continue;
95                    }
96
97                    let name = block
98                        .get("tool_use_id")
99                        .and_then(|v| v.as_str())
100                        .and_then(|id| tool_names.get(id))
101                        .map(String::as_str);
102                    let kind = name.map_or(ToolResultKind::Other, tool_kind::classify_tool_name);
103
104                    if let Some(inner_content) = block.get_mut("content") {
105                        modified |= compress_content_field(inner_content, name, kind);
106                    }
107                }
108            }
109        }
110
111        // Frozen-region user prose: free-text `text` blocks of user turns in
112        // `[cached, boundary)`. Cache-safe by construction — the cached prefix
113        // and the live tail (`>= boundary`) are both left intact, and the
114        // rewrite is content-deterministic so the prefix stays byte-stable.
115        if let Some(a) = user_aggr {
116            let end = boundary.min(messages.len());
117            let start = cached.min(end);
118            for msg in &mut messages[start..end] {
119                if msg.get("role").and_then(|r| r.as_str()) == Some("user")
120                    && let Some(content) = msg.get_mut("content").and_then(|c| c.as_array_mut())
121                {
122                    prose_segments += u64::from(prose::compress_text_blocks(content, a));
123                }
124            }
125        }
126    }
127
128    if prose_segments > 0 {
129        modified = true;
130    }
131    // Every rewrite above lands strictly inside the cache-safe frozen window,
132    // so report the activity as cache-safe (the production invariant gauge).
133    cache_safety::record(prose_segments, true);
134
135    let out = serde_json::to_vec(&doc).unwrap_or_default();
136    let compressed_size = if modified { out.len() } else { original_size };
137    (out, original_size, compressed_size)
138}
139
140/// Compresses a tool_result `content` field unless it is a protected file/source
141/// read, which must reach the model intact (it is what gets edited).
142fn compress_content_field(
143    content: &mut Value,
144    tool_name: Option<&str>,
145    kind: ToolResultKind,
146) -> bool {
147    match content {
148        Value::String(s) => {
149            if should_protect(kind, s) {
150                return false;
151            }
152            let compressed = compress_tool_result(s, tool_name);
153            if compressed.len() < s.len() {
154                *s = compressed;
155                return true;
156            }
157            false
158        }
159        Value::Array(arr) => {
160            let mut modified = false;
161            for item in arr.iter_mut() {
162                if item.get("type").and_then(|t| t.as_str()) == Some("text")
163                    && let Some(text) = item
164                        .get_mut("text")
165                        .and_then(|t| t.as_str().map(String::from))
166                {
167                    if should_protect(kind, &text) {
168                        continue;
169                    }
170                    let compressed = compress_tool_result(&text, tool_name);
171                    if compressed.len() < text.len() {
172                        item["text"] = Value::String(compressed);
173                        modified = true;
174                    }
175                }
176            }
177            modified
178        }
179        _ => false,
180    }
181}
182
183#[cfg(test)]
184mod tests {
185    use super::*;
186
187    fn source_file_body() -> Vec<u8> {
188        let code = (0..60)
189            .map(|i| format!("    let binding_{i} = compute_value_{i}(context, options);"))
190            .collect::<Vec<_>>()
191            .join("\n");
192        let body = serde_json::json!({
193            "model": "claude-opus-4-8",
194            "messages": [
195                {
196                    "role": "assistant",
197                    "content": [{"type": "tool_use", "id": "toolu_1", "name": "Read", "input": {"file_path": "src/app.rs"}}]
198                },
199                {
200                    "role": "user",
201                    "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": code}]
202                }
203            ]
204        });
205        serde_json::to_vec(&body).unwrap()
206    }
207
208    #[test]
209    fn read_tool_result_is_never_truncated() {
210        let bytes = source_file_body();
211        let body: Value = serde_json::from_slice(&bytes).unwrap();
212        let (out, _orig, _comp) = compress_request_body(body, bytes.len());
213        let parsed: Value = serde_json::from_slice(&out).unwrap();
214        let content = parsed["messages"][1]["content"][0]["content"]
215            .as_str()
216            .unwrap();
217        assert!(
218            content.contains("binding_59"),
219            "the full source body must survive — refactors need it intact"
220        );
221        assert!(!content.contains("lines omitted"));
222    }
223
224    fn forge_log_body(tool_name: &str) -> Value {
225        // Generic, highly-repetitive log with no `$ cmd` hint, so routing falls
226        // back to the tool name (exercising the foreign-tool classification)
227        // and the generic compressor (not a command-specific pattern).
228        let mut log = String::new();
229        for i in 0..90 {
230            log.push_str(&format!(
231                "INFO  processing item {i}: ok, latency={i}ms, queue depth normal, retries 0\n"
232            ));
233        }
234        serde_json::json!({
235            "messages": [
236                {"role": "assistant", "content": [{"type": "tool_use", "id": "f1", "name": tool_name, "input": {}}]},
237                {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "f1", "content": log}]}
238            ]
239        })
240    }
241
242    #[test]
243    fn forge_shell_tool_result_compresses() {
244        // A vendor-prefixed foreign shell tool reaches the proxy; its log output
245        // must still be compressed (rtk/ctx_* never see another server's tools).
246        let body = forge_log_body("forge_shell");
247        let bytes = serde_json::to_vec(&body).unwrap();
248        let (_out, orig, comp) = compress_request_body(body, bytes.len());
249        assert!(comp < orig, "foreign shell output must be compressed");
250    }
251
252    #[test]
253    fn foreign_read_tool_protects_source() {
254        // `forge_read` is classified FileRead via the segment fallback, so the
255        // source body must reach the model intact (it is what gets edited).
256        let code = (0..60)
257            .map(|i| format!("    let binding_{i} = compute_value_{i}(context, options);"))
258            .collect::<Vec<_>>()
259            .join("\n");
260        let body = serde_json::json!({
261            "messages": [
262                {"role": "assistant", "content": [{"type": "tool_use", "id": "r1", "name": "forge_read", "input": {"path": "src/app.rs"}}]},
263                {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "r1", "content": code}]}
264            ]
265        });
266        let bytes = serde_json::to_vec(&body).unwrap();
267        let (out, _orig, _comp) = compress_request_body(body, bytes.len());
268        let parsed: Value = serde_json::from_slice(&out).unwrap();
269        let content = parsed["messages"][1]["content"][0]["content"]
270            .as_str()
271            .unwrap();
272        assert!(
273            content.contains("binding_59"),
274            "source body must survive intact"
275        );
276    }
277
278    #[test]
279    fn compress_request_body_is_deterministic() {
280        // #498: the proxy rewrite must be a pure function of the body so the
281        // provider prompt-cache prefix stays byte-identical across turns.
282        let bytes = serde_json::to_vec(&forge_log_body("Bash")).unwrap();
283        let a = compress_request_body(serde_json::from_slice(&bytes).unwrap(), bytes.len()).0;
284        let b = compress_request_body(serde_json::from_slice(&bytes).unwrap(), bytes.len()).0;
285        assert_eq!(a, b, "identical input must yield byte-identical output");
286    }
287
288    /// Long, duplicate-rich natural-language prose that compresses cleanly.
289    fn big_prose() -> String {
290        let p = "You are a careful, senior software engineer. You always explain your \
291                 reasoning before making changes, you prefer small reviewable diffs, and \
292                 you never introduce mock data or placeholders into production code. ";
293        [p; 6].join("\n")
294    }
295
296    #[test]
297    fn system_prose_compressed_and_assistant_untouched() {
298        let _iso = crate::core::data_dir::isolated_data_dir();
299        crate::core::config::Config::update_global(|c| {
300            c.proxy.role_aggressiveness.system = Some(0.6);
301            c.proxy.role_aggressiveness.user = Some(0.6);
302        })
303        .unwrap();
304
305        let prose = big_prose();
306        let assistant_text = big_prose();
307        let body = serde_json::json!({
308            "model": "claude-opus-4-8",
309            "system": prose,
310            "messages": [
311                {"role": "user", "content": [{"type": "text", "text": prose}]},
312                {"role": "assistant", "content": assistant_text},
313            ]
314        });
315        let bytes = serde_json::to_vec(&body).unwrap();
316        let (out, _orig, _comp) = compress_request_body(body, bytes.len());
317        let parsed: Value = serde_json::from_slice(&out).unwrap();
318
319        assert!(
320            parsed["system"].as_str().unwrap().len() < prose.len(),
321            "system prose must be compressed when enabled"
322        );
323        assert_eq!(
324            parsed["messages"][1]["content"].as_str().unwrap(),
325            assistant_text,
326            "assistant turns must pass through verbatim (#710)"
327        );
328    }
329
330    #[test]
331    fn user_prose_compressed_only_in_frozen_region() {
332        let _iso = crate::core::data_dir::isolated_data_dir();
333        crate::core::config::Config::update_global(|c| {
334            c.proxy.role_aggressiveness.user = Some(0.7);
335        })
336        .unwrap();
337
338        let prose = big_prose();
339        // 30 messages → cache-aware boundary = ((30-8)/16)*16 = 16.
340        let mut messages = Vec::new();
341        for i in 0..30 {
342            let role = if i % 2 == 0 { "user" } else { "assistant" };
343            messages.push(serde_json::json!({
344                "role": role,
345                "content": [{"type": "text", "text": prose}]
346            }));
347        }
348        let body = serde_json::json!({ "messages": messages });
349        let bytes = serde_json::to_vec(&body).unwrap();
350        let (out, _o, _c) = compress_request_body(body, bytes.len());
351        let parsed: Value = serde_json::from_slice(&out).unwrap();
352
353        let frozen_user = parsed["messages"][0]["content"][0]["text"]
354            .as_str()
355            .unwrap();
356        assert!(
357            frozen_user.len() < prose.len(),
358            "user prose in the frozen region must be compressed"
359        );
360        assert_eq!(
361            parsed["messages"][1]["content"][0]["text"]
362                .as_str()
363                .unwrap(),
364            prose,
365            "assistant prose is never compressed"
366        );
367        let live_tail_user = parsed["messages"][28]["content"][0]["text"]
368            .as_str()
369            .unwrap();
370        assert_eq!(
371            live_tail_user, prose,
372            "user prose in the live tail (>= boundary) must be preserved for quality"
373        );
374    }
375
376    #[test]
377    fn client_cached_prefix_disables_system_prose() {
378        let _iso = crate::core::data_dir::isolated_data_dir();
379        crate::core::config::Config::update_global(|c| {
380            c.proxy.role_aggressiveness.system = Some(0.9);
381        })
382        .unwrap();
383
384        let prose = big_prose();
385        let body = serde_json::json!({
386            "system": prose,
387            "messages": [
388                {"role": "user", "content": [
389                    {"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}
390                ]},
391                {"role": "assistant", "content": "ok"}
392            ]
393        });
394        let bytes = serde_json::to_vec(&body).unwrap();
395        let (out, _o, _c) = compress_request_body(body, bytes.len());
396        let parsed: Value = serde_json::from_slice(&out).unwrap();
397        assert_eq!(
398            parsed["system"].as_str().unwrap(),
399            prose,
400            "system must stay verbatim when the client caches a message prefix (#448)"
401        );
402    }
403
404    #[test]
405    fn prose_compression_is_deterministic() {
406        let _iso = crate::core::data_dir::isolated_data_dir();
407        crate::core::config::Config::update_global(|c| {
408            c.proxy.role_aggressiveness.system = Some(0.6);
409        })
410        .unwrap();
411        let prose = big_prose();
412        let mk = || serde_json::json!({"system": prose, "messages": [{"role": "user", "content": "hi"}]});
413        let (a, b) = (mk(), mk());
414        let la = serde_json::to_vec(&a).unwrap().len();
415        let lb = serde_json::to_vec(&b).unwrap().len();
416        assert_eq!(
417            compress_request_body(a, la).0,
418            compress_request_body(b, lb).0,
419            "prose compression must be byte-identical for identical input (#498)"
420        );
421    }
422
423    #[test]
424    fn bash_tool_result_still_compresses() {
425        let log = {
426            let mut s = String::from(
427                "$ 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",
428            );
429            for i in 0..90 {
430                s.push_str(&format!("\tmodified:   src/module_{i}/file_{i}.rs\n"));
431            }
432            s.push_str("\nno changes added to commit (use \"git add\" and/or \"git commit -a\")\n");
433            s
434        };
435        let body = serde_json::json!({
436            "messages": [
437                {"role": "assistant", "content": [{"type": "tool_use", "id": "t1", "name": "Bash", "input": {}}]},
438                {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "t1", "content": log}]}
439            ]
440        });
441        let bytes = serde_json::to_vec(&body).unwrap();
442        let (_out, orig, comp) = compress_request_body(body, bytes.len());
443        assert!(comp < orig, "shell output must still be compressed");
444    }
445}