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::compress::compress_tool_result;
10use super::forward;
11use super::tool_kind::{self, should_protect, ToolResultKind};
12use super::ProxyState;
13
14/// Conversation turns kept fully intact at the tail of the history; older
15/// tool results are summarized by `history_prune`.
16const KEEP_RECENT: usize = 6;
17
18pub async fn handler(
19    State(state): State<ProxyState>,
20    req: Request<Body>,
21) -> Result<Response, StatusCode> {
22    let upstream = state.anthropic_upstream.clone();
23    forward::forward_request(
24        State(state),
25        req,
26        &upstream,
27        "/v1/messages",
28        compress_request_body,
29        "Anthropic",
30        &[],
31    )
32    .await
33}
34
35fn compress_request_body(body: &[u8]) -> (Vec<u8>, usize, usize) {
36    let original_size = body.len();
37
38    let parsed: Value = match serde_json::from_slice(body) {
39        Ok(v) => v,
40        Err(_) => return (body.to_vec(), original_size, original_size),
41    };
42
43    let mut doc = parsed;
44    let mut modified = false;
45
46    if let Some(messages) = doc.get_mut("messages").and_then(|m| m.as_array_mut()) {
47        // Resolve tool-call id → tool name so file/source reads can be protected
48        // from lossy compression that would force the model to re-read mid-task.
49        let tool_names = tool_kind::anthropic_tool_names(messages);
50
51        super::history_prune::prune_history(messages, KEEP_RECENT, &tool_names);
52        modified = true;
53
54        for msg in messages.iter_mut() {
55            let role = msg.get("role").and_then(|r| r.as_str()).unwrap_or("");
56            if role != "user" {
57                continue;
58            }
59
60            if let Some(content) = msg.get_mut("content").and_then(|c| c.as_array_mut()) {
61                for block in content.iter_mut() {
62                    if block.get("type").and_then(|t| t.as_str()) != Some("tool_result") {
63                        continue;
64                    }
65
66                    let name = block
67                        .get("tool_use_id")
68                        .and_then(|v| v.as_str())
69                        .and_then(|id| tool_names.get(id))
70                        .map(String::as_str);
71                    let kind = name.map_or(ToolResultKind::Other, tool_kind::classify_tool_name);
72
73                    if let Some(inner_content) = block.get_mut("content") {
74                        modified |= compress_content_field(inner_content, name, kind);
75                    }
76                }
77            }
78        }
79    }
80
81    if !modified {
82        return (body.to_vec(), original_size, original_size);
83    }
84
85    match serde_json::to_vec(&doc) {
86        Ok(compressed) => {
87            let compressed_size = compressed.len();
88            (compressed, original_size, compressed_size)
89        }
90        Err(_) => (body.to_vec(), original_size, original_size),
91    }
92}
93
94/// Compresses a tool_result `content` field unless it is a protected file/source
95/// read, which must reach the model intact (it is what gets edited).
96fn compress_content_field(
97    content: &mut Value,
98    tool_name: Option<&str>,
99    kind: ToolResultKind,
100) -> bool {
101    match content {
102        Value::String(s) => {
103            if should_protect(kind, s) {
104                return false;
105            }
106            let compressed = compress_tool_result(s, tool_name);
107            if compressed.len() < s.len() {
108                *s = compressed;
109                return true;
110            }
111            false
112        }
113        Value::Array(arr) => {
114            let mut modified = false;
115            for item in arr.iter_mut() {
116                if item.get("type").and_then(|t| t.as_str()) == Some("text") {
117                    if let Some(text) = item
118                        .get_mut("text")
119                        .and_then(|t| t.as_str().map(String::from))
120                    {
121                        if should_protect(kind, &text) {
122                            continue;
123                        }
124                        let compressed = compress_tool_result(&text, tool_name);
125                        if compressed.len() < text.len() {
126                            item["text"] = Value::String(compressed);
127                            modified = true;
128                        }
129                    }
130                }
131            }
132            modified
133        }
134        _ => false,
135    }
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141
142    fn source_file_body() -> Vec<u8> {
143        let code = (0..60)
144            .map(|i| format!("    let binding_{i} = compute_value_{i}(context, options);"))
145            .collect::<Vec<_>>()
146            .join("\n");
147        let body = serde_json::json!({
148            "model": "claude-opus-4-8",
149            "messages": [
150                {
151                    "role": "assistant",
152                    "content": [{"type": "tool_use", "id": "toolu_1", "name": "Read", "input": {"file_path": "src/app.rs"}}]
153                },
154                {
155                    "role": "user",
156                    "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": code}]
157                }
158            ]
159        });
160        serde_json::to_vec(&body).unwrap()
161    }
162
163    #[test]
164    fn read_tool_result_is_never_truncated() {
165        let bytes = source_file_body();
166        let (out, _orig, _comp) = compress_request_body(&bytes);
167        let parsed: Value = serde_json::from_slice(&out).unwrap();
168        let content = parsed["messages"][1]["content"][0]["content"]
169            .as_str()
170            .unwrap();
171        assert!(
172            content.contains("binding_59"),
173            "the full source body must survive — refactors need it intact"
174        );
175        assert!(!content.contains("lines omitted"));
176    }
177
178    #[test]
179    fn bash_tool_result_still_compresses() {
180        let log = {
181            let mut s = String::from(
182                "$ 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",
183            );
184            for i in 0..90 {
185                s.push_str(&format!("\tmodified:   src/module_{i}/file_{i}.rs\n"));
186            }
187            s.push_str("\nno changes added to commit (use \"git add\" and/or \"git commit -a\")\n");
188            s
189        };
190        let body = serde_json::json!({
191            "messages": [
192                {"role": "assistant", "content": [{"type": "tool_use", "id": "t1", "name": "Bash", "input": {}}]},
193                {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "t1", "content": log}]}
194            ]
195        });
196        let bytes = serde_json::to_vec(&body).unwrap();
197        let (_out, orig, comp) = compress_request_body(&bytes);
198        assert!(comp < orig, "shell output must still be compressed");
199    }
200}