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};
13
14pub async fn handler(
15    State(state): State<ProxyState>,
16    req: Request<Body>,
17) -> Result<Response, StatusCode> {
18    let upstream = state.anthropic_upstream();
19    forward::forward_request(
20        State(state),
21        req,
22        &upstream,
23        "/v1/messages",
24        compress_request_body,
25        "Anthropic",
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        // Resolve tool-call id → tool name so file/source reads can be protected
37        // from lossy compression that would force the model to re-read mid-task.
38        let tool_names = tool_kind::anthropic_tool_names(messages);
39
40        // Prune at a frozen, cache-aware boundary by default: Anthropic's
41        // prompt cache matches exact prefixes, so the boundary must not move
42        // every turn (see `history_prune::prune_boundary`).
43        let mode = crate::core::config::Config::load()
44            .proxy
45            .resolved_history_mode();
46        let boundary = super::history_prune::prune_boundary(mode, messages.len());
47        // Never rewrite content the client has marked with `cache_control`:
48        // pruning inside the already-cached prefix invalidates Anthropic's
49        // prompt cache from the first changed message (#448). Pruning therefore
50        // starts after the last breakpoint; with no breakpoint this is 0, i.e.
51        // the previous behaviour.
52        let cached = super::history_prune::cached_prefix_len(messages);
53        modified |=
54            super::history_prune::prune_history_range(messages, cached, boundary, &tool_names);
55
56        for msg in messages.iter_mut() {
57            let role = msg.get("role").and_then(|r| r.as_str()).unwrap_or("");
58            if role != "user" {
59                continue;
60            }
61
62            if let Some(content) = msg.get_mut("content").and_then(|c| c.as_array_mut()) {
63                for block in content.iter_mut() {
64                    if block.get("type").and_then(|t| t.as_str()) != Some("tool_result") {
65                        continue;
66                    }
67
68                    let name = block
69                        .get("tool_use_id")
70                        .and_then(|v| v.as_str())
71                        .and_then(|id| tool_names.get(id))
72                        .map(String::as_str);
73                    let kind = name.map_or(ToolResultKind::Other, tool_kind::classify_tool_name);
74
75                    if let Some(inner_content) = block.get_mut("content") {
76                        modified |= compress_content_field(inner_content, name, kind);
77                    }
78                }
79            }
80        }
81    }
82
83    let out = serde_json::to_vec(&doc).unwrap_or_default();
84    let compressed_size = if modified { out.len() } else { original_size };
85    (out, original_size, compressed_size)
86}
87
88/// Compresses a tool_result `content` field unless it is a protected file/source
89/// read, which must reach the model intact (it is what gets edited).
90fn compress_content_field(
91    content: &mut Value,
92    tool_name: Option<&str>,
93    kind: ToolResultKind,
94) -> bool {
95    match content {
96        Value::String(s) => {
97            if should_protect(kind, s) {
98                return false;
99            }
100            let compressed = compress_tool_result(s, tool_name);
101            if compressed.len() < s.len() {
102                *s = compressed;
103                return true;
104            }
105            false
106        }
107        Value::Array(arr) => {
108            let mut modified = false;
109            for item in arr.iter_mut() {
110                if item.get("type").and_then(|t| t.as_str()) == Some("text")
111                    && let Some(text) = item
112                        .get_mut("text")
113                        .and_then(|t| t.as_str().map(String::from))
114                {
115                    if should_protect(kind, &text) {
116                        continue;
117                    }
118                    let compressed = compress_tool_result(&text, tool_name);
119                    if compressed.len() < text.len() {
120                        item["text"] = Value::String(compressed);
121                        modified = true;
122                    }
123                }
124            }
125            modified
126        }
127        _ => false,
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use super::*;
134
135    fn source_file_body() -> Vec<u8> {
136        let code = (0..60)
137            .map(|i| format!("    let binding_{i} = compute_value_{i}(context, options);"))
138            .collect::<Vec<_>>()
139            .join("\n");
140        let body = serde_json::json!({
141            "model": "claude-opus-4-8",
142            "messages": [
143                {
144                    "role": "assistant",
145                    "content": [{"type": "tool_use", "id": "toolu_1", "name": "Read", "input": {"file_path": "src/app.rs"}}]
146                },
147                {
148                    "role": "user",
149                    "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": code}]
150                }
151            ]
152        });
153        serde_json::to_vec(&body).unwrap()
154    }
155
156    #[test]
157    fn read_tool_result_is_never_truncated() {
158        let bytes = source_file_body();
159        let body: Value = serde_json::from_slice(&bytes).unwrap();
160        let (out, _orig, _comp) = compress_request_body(body, bytes.len());
161        let parsed: Value = serde_json::from_slice(&out).unwrap();
162        let content = parsed["messages"][1]["content"][0]["content"]
163            .as_str()
164            .unwrap();
165        assert!(
166            content.contains("binding_59"),
167            "the full source body must survive — refactors need it intact"
168        );
169        assert!(!content.contains("lines omitted"));
170    }
171
172    fn forge_log_body(tool_name: &str) -> Value {
173        // Generic, highly-repetitive log with no `$ cmd` hint, so routing falls
174        // back to the tool name (exercising the foreign-tool classification)
175        // and the generic compressor (not a command-specific pattern).
176        let mut log = String::new();
177        for i in 0..90 {
178            log.push_str(&format!(
179                "INFO  processing item {i}: ok, latency={i}ms, queue depth normal, retries 0\n"
180            ));
181        }
182        serde_json::json!({
183            "messages": [
184                {"role": "assistant", "content": [{"type": "tool_use", "id": "f1", "name": tool_name, "input": {}}]},
185                {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "f1", "content": log}]}
186            ]
187        })
188    }
189
190    #[test]
191    fn forge_shell_tool_result_compresses() {
192        // A vendor-prefixed foreign shell tool reaches the proxy; its log output
193        // must still be compressed (rtk/ctx_* never see another server's tools).
194        let body = forge_log_body("forge_shell");
195        let bytes = serde_json::to_vec(&body).unwrap();
196        let (_out, orig, comp) = compress_request_body(body, bytes.len());
197        assert!(comp < orig, "foreign shell output must be compressed");
198    }
199
200    #[test]
201    fn foreign_read_tool_protects_source() {
202        // `forge_read` is classified FileRead via the segment fallback, so the
203        // source body must reach the model intact (it is what gets edited).
204        let code = (0..60)
205            .map(|i| format!("    let binding_{i} = compute_value_{i}(context, options);"))
206            .collect::<Vec<_>>()
207            .join("\n");
208        let body = serde_json::json!({
209            "messages": [
210                {"role": "assistant", "content": [{"type": "tool_use", "id": "r1", "name": "forge_read", "input": {"path": "src/app.rs"}}]},
211                {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "r1", "content": code}]}
212            ]
213        });
214        let bytes = serde_json::to_vec(&body).unwrap();
215        let (out, _orig, _comp) = compress_request_body(body, bytes.len());
216        let parsed: Value = serde_json::from_slice(&out).unwrap();
217        let content = parsed["messages"][1]["content"][0]["content"]
218            .as_str()
219            .unwrap();
220        assert!(
221            content.contains("binding_59"),
222            "source body must survive intact"
223        );
224    }
225
226    #[test]
227    fn compress_request_body_is_deterministic() {
228        // #498: the proxy rewrite must be a pure function of the body so the
229        // provider prompt-cache prefix stays byte-identical across turns.
230        let bytes = serde_json::to_vec(&forge_log_body("Bash")).unwrap();
231        let a = compress_request_body(serde_json::from_slice(&bytes).unwrap(), bytes.len()).0;
232        let b = compress_request_body(serde_json::from_slice(&bytes).unwrap(), bytes.len()).0;
233        assert_eq!(a, b, "identical input must yield byte-identical output");
234    }
235
236    #[test]
237    fn bash_tool_result_still_compresses() {
238        let log = {
239            let mut s = String::from(
240                "$ 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",
241            );
242            for i in 0..90 {
243                s.push_str(&format!("\tmodified:   src/module_{i}/file_{i}.rs\n"));
244            }
245            s.push_str("\nno changes added to commit (use \"git add\" and/or \"git commit -a\")\n");
246            s
247        };
248        let body = serde_json::json!({
249            "messages": [
250                {"role": "assistant", "content": [{"type": "tool_use", "id": "t1", "name": "Bash", "input": {}}]},
251                {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "t1", "content": log}]}
252            ]
253        });
254        let bytes = serde_json::to_vec(&body).unwrap();
255        let (_out, orig, comp) = compress_request_body(body, bytes.len());
256        assert!(comp < orig, "shell output must still be compressed");
257    }
258}