1. 1#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
2. 2
3. 3use tokio_util::sync::CancellationToken;
4. 4
5. 5use vak_llm::Provider;
6. 6use vak_llm::anthropic::{AnthropicConfig, AnthropicProvider, build_body, map_status_error};
7. 7use vak_llm::error::LlmError;
8. 8use vak_llm::sse::SseDecoder;
9. 9use vak_llm::stream::StreamEvent;
10. 10use vak_llm::types::{ChatRequest, ContentBlock, Message, Role, StopReason, ToolDefinition};
11. 11
12. 12fn sample_request() -> ChatRequest {
13. 13 let mut req = ChatRequest::new("claude-sonnet-4-5");
14. 14 req.system = Some("You are a coding agent.".into());
15. 15 req.messages = vec![
16. 16 Message::user_text("list the files"),
17. 17 Message::assistant(vec![ContentBlock::ToolUse {
18. 18 id: "toolu_1".into(),
19. 19 name: "bash".into(),
20. 20 input: serde_json::json!({"command": "ls"}),
21. 21 }]),
22. 22 Message {
23. 23 role: Role::User,
24. 24 content: vec![ContentBlock::tool_result("toolu_1", "src\nREADME.md")],
25. 25 },
26. 26 ];
27. 27 req.tools = vec![ToolDefinition::new(
28. 28 "bash",
29. 29 "Run a shell command",
30. 30 serde_json::json!({"type": "object", "properties": {"command": {"type": "string"}}, "required": ["command"]}),
31. 31 )];
32. 32 req
33. 33}
34. 34
35. 35#[test]
36. 36fn body_serialization_matches_anthropic_shape() {
37. 37 let body = build_body(&sample_request(), true, false).unwrap();
38. 38 assert_eq!(body["model"], "claude-sonnet-4-5");
39. 39 assert_eq!(body["max_tokens"], 8192);
40. 40 assert_eq!(body["stream"], true);
41. 41 let system = body["system"].as_array().unwrap();
42. 42 assert_eq!(system[0]["type"], "text");
43. 43 assert_eq!(system[0]["text"], "You are a coding agent.");
44. 44 assert_eq!(system[0]["cache_control"]["type"], "ephemeral");
45. 45 let tools = body["tools"].as_array().unwrap();
46. 46 assert_eq!(tools[0]["name"], "bash");
47. 47 assert!(tools[0]["input_schema"].is_object());
48. 48 let messages = body["messages"].as_array().unwrap();
49. 49 assert_eq!(messages.len(), 3);
50. 50 assert_eq!(messages[2]["content"][0]["type"], "tool_result");
51. 51 assert_eq!(messages[2]["content"][0]["tool_use_id"], "toolu_1");
52. 52}
53. 53
54. 54#[test]
55. 55fn body_rejects_misplaced_blocks() {
56. 56 let mut req = ChatRequest::new("m");
57. 57 req.messages = vec![Message {
58. 58 role: Role::Assistant,
59. 59 content: vec![ContentBlock::tool_result("x", "y")],
60. 60 }];
61. 61 assert!(matches!(
62. 62 build_body(&req, true, false),
63. 63 Err(LlmError::InvalidRequest(_))
64. 64 ));
65. 65}
66. 66
67. 67#[test]
68. 68fn status_error_mapping() {
69. 69 assert!(matches!(
70. 70 map_status_error(401, r#"{"error":{"message":"bad key"}}"#, None),
71. 71 LlmError::Auth(_)
72. 72 ));
73. 73 match map_status_error(429, r#"{"error":{"message":"slow down"}}"#, Some(7)) {
74. 74 LlmError::RateLimit {
75. 75 retry_after_secs, ..
76. 76 } => assert_eq!(retry_after_secs, Some(7)),
77. 77 other => panic!("expected rate limit, got {other:?}"),
78. 78 }
79. 79 assert!(matches!(
80. 80 map_status_error(529, r#"{"error":{"message":"overloaded"}}"#, None),
81. 81 LlmError::Overloaded(_)
82. 82 ));
83. 83}
84. 84
85. 85const FIXTURE_STREAM: &str = "\
86. 86event: message_start\n\
87. 87data: {\"type\":\"message_start\",\"message\":{\"model\":\"claude-sonnet-4-5\",\"usage\":{\"input_tokens\":10,\"output_tokens\":0}}}\n\
88. 88\n\
89. 89event: content_block_start\n\
90. 90data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\"}}\n\
91. 91\n\
92. 92event: content_block_delta\n\
93. 93data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Hel\"}}\n\
94. 94\n\
95. 95event: content_block_delta\n\
96. 96data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"lo\"}}\n\
97. 97\n\
98. 98event: content_block_stop\n\
99. 99data: {\"type\":\"content_block_stop\",\"index\":0}\n\
100. 100\n\
101. 101event: content_block_start\n\
102. 102data: {\"type\":\"content_block_start\",\"index\":1,\"content_block\":{\"type\":\"tool_use\",\"id\":\"toolu_9\",\"name\":\"read\"}}\n\
103. 103\n\
104. 104event: content_block_delta\n\
105. 105data: {\"type\":\"content_block_delta\",\"index\":1,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"path\\\":\"}}\n\
106. 106\n\
107. 107event: content_block_delta\n\
108. 108data: {\"type\":\"content_block_delta\",\"index\":1,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"\\\"a.txt\\\"}\"}}\n\
109. 109\n\
110. 110event: message_delta\n\
111. 111data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"output_tokens\":42}}\n\
112. 112\n\
113. 113event: message_stop\n\
114. 114data: {\"type\":\"message_stop\"}\n\
115. 115\n";
116. 116
117. 117#[tokio::test]
118. 118async fn full_stream_conversion_accumulates_snapshot() {
119. 119 let provider = AnthropicProvider::new(AnthropicConfig {
120. 120 api_key: "k".into(),
121. 121 base_url: mock_server_url(FIXTURE_STREAM).await,
122. 122 model: String::new(),
123. 123 fast_mode: false,
124. 124 })
125. 125 .unwrap();
126. 126
127. 127 let cancel = CancellationToken::new();
128. 128 let mut es = provider.stream(sample_request(), cancel).await.unwrap();
129. 129 let mut text = String::new();
130. 130 let mut snapshots_consistent = true;
131. 131 while let Some(ev) = futures::StreamExt::next(&mut es).await {
132. 132 if let StreamEvent::TextDelta { delta, partial } = &ev {
133. 133 text.push_str(delta);
134. 134 let joined: String = partial
135. 135 .content
136. 136 .iter()
137. 137 .filter_map(|b| match b {
138. 138 ContentBlock::Text { text } => Some(text.as_str()),
139. 139 _ => None,
140. 140 })
141. 141 .collect();
142. 142 if joined != text {
143. 143 snapshots_consistent = false;
144. 144 }
145. 145 }
146. 146 }
147. 147 let final_msg = es.result().await.unwrap();
148. 148 assert!(snapshots_consistent);
149. 149 assert_eq!(text, "Hello");
150. 150 assert_eq!(final_msg.stop_reason, StopReason::ToolUse);
151. 151 assert_eq!(final_msg.usage.output_tokens, 42);
152. 152 assert_eq!(final_msg.usage.input_tokens, 10);
153. 153 let tool_calls: Vec<&ContentBlock> = final_msg
154. 154 .content
155. 155 .iter()
156. 156 .filter(|b| matches!(b, ContentBlock::ToolUse { .. }))
157. 157 .collect();
158. 158 assert_eq!(tool_calls.len(), 1);
159. 159 if let ContentBlock::ToolUse { id, name, input } = tool_calls[0] {
160. 160 assert_eq!(id, "toolu_9");
161. 161 assert_eq!(name, "read");
162. 162 assert_eq!(input, &serde_json::json!({"path": "a.txt"}));
163. 163 } else {
164. 164 unreachable!()
165. 165 }
166. 166}
167. 167
168. 168// Realistic cache-bearing usage (docs/design/68-context-engine.md §1):
169. 169// Anthropic already reports `input_tokens` as the non-cached remainder,
170. 170// with `cache_read_input_tokens`/`cache_creation_input_tokens` carried
171. 171// alongside it — this fixture pins that shape, and that `message_delta`
172. 172// (which only ever carries `output_tokens`) never disturbs it.
173. 173const FIXTURE_CACHED_USAGE: &str = "\
174. 174event: message_start\n\
175. 175data: {\"type\":\"message_start\",\"message\":{\"model\":\"claude-sonnet-4-5\",\"usage\":{\"input_tokens\":50,\"output_tokens\":0,\"cache_read_input_tokens\":9000,\"cache_creation_input_tokens\":300}}}\n\
176. 176\n\
177. 177event: message_delta\n\
178. 178data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":15}}\n\
179. 179\n\
180. 180event: message_stop\n\
181. 181data: {\"type\":\"message_stop\"}\n\
182. 182\n";
183. 183
184. 184#[tokio::test]
185. 185async fn cache_read_and_creation_tokens_survive_message_delta_untouched() {
186. 186 let provider = AnthropicProvider::new(AnthropicConfig {
187. 187 api_key: "k".into(),
188. 188 base_url: mock_server_url(FIXTURE_CACHED_USAGE).await,
189. 189 model: String::new(),
190. 190 fast_mode: false,
191. 191 })
192. 192 .unwrap();
193. 193 let mut es = provider
194. 194 .stream(ChatRequest::new("m"), CancellationToken::new())
195. 195 .await
196. 196 .unwrap();
197. 197 while futures::StreamExt::next(&mut es).await.is_some() {}
198. 198 let msg = es.result().await.unwrap();
199. 199
200. 200 assert_eq!(msg.usage.input_tokens, 50);
201. 201 assert_eq!(msg.usage.cache_read_input_tokens, Some(9_000));
202. 202 assert_eq!(msg.usage.cache_creation_input_tokens, Some(300));
203. 203 // message_delta only ever carries output_tokens on this wire; the
204. 204 // input/cache figures set at message_start must be untouched by it.
205. 205 assert_eq!(msg.usage.output_tokens, 15);
206. 206 assert_eq!(msg.usage.prompt_tokens(), 9_350);
207. 207}
208. 208
209. 209#[tokio::test]
210. 210async fn abort_preserves_partial_output() {
211. 211 let part1 = "\
212. 212event: message_start\n\
213. 213data: {\"type\":\"message_start\",\"message\":{\"model\":\"m\",\"usage\":{}}}\n\
214. 214\n\
215. 215event: content_block_start\n\
216. 216data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\"}}\n\
217. 217\n\
218. 218event: content_block_delta\n\
219. 219data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"partial answer\"}}\n\
220. 220\n"
221. 221 .to_string();
222. 222 let url = delayed_server_url(part1, String::new()).await;
223. 223 let provider = AnthropicProvider::new(AnthropicConfig {
224. 224 api_key: "k".into(),
225. 225 base_url: url,
226. 226 model: String::new(),
227. 227 fast_mode: false,
228. 228 })
229. 229 .unwrap();
230. 230
231. 231 let cancel = CancellationToken::new();
232. 232 let mut es = provider
233. 233 .stream(sample_request(), cancel.clone())
234. 234 .await
235. 235 .unwrap();
236. 236 while let Some(ev) = futures::StreamExt::next(&mut es).await {
237. 237 if matches!(ev, StreamEvent::TextDelta { .. }) {
238. 238 break;
239. 239 }
240. 240 }
241. 241 cancel.cancel();
242. 242 match es.result().await {
243. 243 Err(LlmError::Aborted { partial }) => {
244. 244 let p = partial.expect("partial should be preserved");
245. 245 assert_eq!(
246. 246 p.content[0],
247. 247 ContentBlock::Text {
248. 248 text: "partial answer".into()
249. 249 }
250. 250 );
251. 251 }
252. 252 other => panic!("expected aborted with partial, got {other:?}"),
253. 253 }
254. 254}
255. 255
256. 256#[tokio::test]
257. 257async fn http_error_maps_to_typed_value() {
258. 258 let url = error_server_url(401, r#"{"error":{"message":"invalid x-api-key"}}"#).await;
259. 259 let provider = AnthropicProvider::new(AnthropicConfig {
260. 260 api_key: "bad".into(),
261. 261 base_url: url,
262. 262 model: String::new(),
263. 263 fast_mode: false,
264. 264 })
265. 265 .unwrap();
266. 266 let err = provider
267. 267 .stream(ChatRequest::new("m"), CancellationToken::new())
268. 268 .await
269. 269 .err()
270. 270 .expect("expected error");
271. 271 assert!(matches!(err, LlmError::Auth(_)), "got {err:?}");
272. 272}
273. 273
274. 274async fn spawn_server<F>(handler: F) -> String
275. 275where
276. 276 F: FnOnce() -> Vec<u8> + Send + 'static,
277. 277{
278. 278 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
279. 279 let addr = listener.local_addr().unwrap();
280. 280 tokio::spawn(async move {
281. 281 let (mut sock, _) = listener.accept().await.unwrap();
282. 282 drain_headers(&mut sock).await;
283. 283 use tokio::io::AsyncWriteExt;
284. 284 let body = handler();
285. 285 let _ = sock.write_all(&body).await;
286. 286 let _ = sock.flush().await;
287. 287 let _ = sock.shutdown().await;
288. 288 });
289. 289 format!("http://{addr}")
290. 290}
291. 291
292. 292async fn drain_headers<S>(sock: &mut S)
293. 293where
294. 294 S: tokio::io::AsyncRead + Unpin,
295. 295{
296. 296 use tokio::io::AsyncReadExt;
297. 297 let mut buf = Vec::new();
298. 298 let mut chunk = [0u8; 1024];
299. 299 loop {
300. 300 let n = sock.read(&mut chunk).await.unwrap_or(0);
301. 301 if n == 0 {
302. 302 return;
303. 303 }
304. 304 buf.extend_from_slice(&chunk[..n]);
305. 305 if buf.windows(4).any(|w| w == b"\r\n\r\n") {
306. 306 return;
307. 307 }
308. 308 }
309. 309}
310. 310
311. 311async fn mock_server_url(sse_body: &str) -> String {
312. 312 let owned = format!(
313. 313 "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\nconnection: close\r\n\r\n{sse_body}"
314. 314 );
315. 315 spawn_server(move || owned.into_bytes()).await
316. 316}
317. 317
318. 318async fn delayed_server_url(first: String, second: String) -> String {
319. 319 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
320. 320 let addr = listener.local_addr().unwrap();
321. 321 tokio::spawn(async move {
322. 322 let (mut sock, _) = listener.accept().await.unwrap();
323. 323 drain_headers(&mut sock).await;
324. 324 use tokio::io::AsyncWriteExt;
325. 325 let head =
326. 326 b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\nconnection: close\r\n\r\n";
327. 327 let _ = sock.write_all(head).await;
328. 328 let _ = sock.write_all(first.as_bytes()).await;
329. 329 let _ = sock.flush().await;
330. 330 tokio::time::sleep(std::time::Duration::from_secs(30)).await;
331. 331 let _ = sock.write_all(second.as_bytes()).await;
332. 332 let _ = sock.shutdown().await;
333. 333 });
334. 334 format!("http://{addr}")
335. 335}
336. 336
337. 337async fn error_server_url(status: u16, json: &str) -> String {
338. 338 let owned = format!(
339. 339 "HTTP/1.1 {status} Error\r\ncontent-type: application/json\r\nconnection: close\r\n\r\n{json}"
340. 340 );
341. 341 spawn_server(move || owned.into_bytes()).await
342. 342}
343. 343
344. 344#[test]
345. 345fn sse_decoder_handles_split_chunks_and_crlf() {
346. 346 let mut d = SseDecoder::new();
347. 347 d.push(b"event: mess");
348. 348 assert!(d.next_frame().is_none());
349. 349 d.push(b"age_start\r\ndata: {\"a\":1}\r\n");
350. 350 d.push(b"\r\ndata: line1\ndata: line2\n\n");
351. 351 let f1 = d.next_frame().unwrap();
352. 352 assert_eq!(f1.event.as_deref(), Some("message_start"));
353. 353 assert_eq!(f1.data, "{\"a\":1}");
354. 354 let f2 = d.next_frame().unwrap();
355. 355 assert_eq!(f2.event, None);
356. 356 assert_eq!(f2.data, "line1\nline2");
357. 357 assert!(d.next_frame().is_none());
358. 358}
359. 359
360. 360#[test]
361. 361fn sse_decoder_skips_comments_and_keepalives() {
362. 362 let mut d = SseDecoder::new();
363. 363 d.push(b": keepalive\n\nevent: x\ndata: y\n\n");
364. 364 let f = d.next_frame().unwrap();
365. 365 assert_eq!(f.data, "y");
366. 366 assert_eq!(f.event.as_deref(), Some("x"));
367. 367}
368. 368