1. 1#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
2. 2
3. 3//! A file sent on a channel comes back edited (docs/design/72, P5): the
4. 4//! Agent drafts a change to the inbox copy with `office_apply`, and the
5. 5//! waiting reply carries the draft, under the name it was sent with and
6. 6//! with a caption saying what changed, to a bridge that accepts files. A
7. 7//! draft with a sensitivity label is held, and a bridge that takes no
8. 8//! files is told where the draft is.
9. 9
10. 10use std::collections::VecDeque;
11. 11use std::sync::{Arc, Mutex};
12. 12
13. 13use base64::Engine as _;
14. 14use sha2::Digest as _;
15. 15use tokio_util::sync::CancellationToken;
16. 16
17. 17use vak_core::Core;
18. 18use vak_llm::stream;
19. 19use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, StopReason, Usage};
20. 20use vak_llm::{EventStream, LlmError, Provider};
21. 21
22. 22struct Scripted(Mutex<VecDeque<AssistantMessage>>, Mutex<Vec<ChatRequest>>);
23. 23
24. 24#[async_trait::async_trait]
25. 25impl Provider for Scripted {
26. 26 fn name(&self) -> &str {
27. 27 "scripted"
28. 28 }
29. 29
30. 30 async fn stream(
31. 31 &self,
32. 32 request: ChatRequest,
33. 33 _cancel: CancellationToken,
34. 34 ) -> Result<EventStream, LlmError> {
35. 35 self.1.lock().unwrap().push(request);
36. 36 // After the script, the last answer again: a side request (a
37. 37 // completion check) must not stall the turn on retries.
38. 38 let next = {
39. 39 let mut script = self.0.lock().unwrap();
40. 40 if script.len() > 1 {
41. 41 script.pop_front()
42. 42 } else {
43. 43 script.front().cloned()
44. 44 }
45. 45 };
46. 46 let (mut sink, rx) = stream::channel(16);
47. 47 match next {
48. 48 Some(message) => {
49. 49 sink.push(stream::StreamEvent::Start {
50. 50 partial: message.clone(),
51. 51 });
52. 52 sink.close_message(message).await;
53. 53 }
54. 54 None => sink.close_error(LlmError::Parse("exhausted".into())).await,
55. 55 }
56. 56 Ok(rx)
57. 57 }
58. 58}
59. 59
60. 60fn message(content: ContentBlock, stop_reason: StopReason) -> AssistantMessage {
61. 61 AssistantMessage {
62. 62 content: vec![content],
63. 63 stop_reason,
64. 64 usage: Usage::default(),
65. 65 model: "test-model".into(),
66. 66 response_id: None,
67. 67 }
68. 68}
69. 69
70. 70/// Where the gateway saves `bytes` sent as `name`, and the file's sha256.
71. 71fn inbox_path(name: &str, bytes: &[u8]) -> (String, String) {
72. 72 let digest = sha2::Sha256::digest(bytes);
73. 73 let hex: String = digest.iter().map(|byte| format!("{byte:02x}")).collect();
74. 74 (format!("inbox/{}-{name}", &hex[..12]), hex)
75. 75}
76. 76
77. 77/// Sends `bytes` as `name` on a channel whose model edits it with `ops`,
78. 78/// and returns the waiting reply.
79. 79async fn round_trip(
80. 80 name: &str,
81. 81 bytes: &[u8],
82. 82 ops: serde_json::Value,
83. 83 accepts_files: bool,
84. 84) -> serde_json::Value {
85. 85 let (path, digest) = inbox_path(name, bytes);
86. 86 let provider = Arc::new(Scripted(
87. 87 Mutex::new(VecDeque::from([
88. 88 message(
89. 89 ContentBlock::ToolUse {
90. 90 id: "call_edit".into(),
91. 91 name: "office_apply".into(),
92. 92 input: serde_json::json!({"path": path, "base_digest": digest, "ops": ops}),
93. 93 },
94. 94 StopReason::ToolUse,
95. 95 ),
96. 96 message(ContentBlock::text("Updated it."), StopReason::EndTurn),
97. 97 ])),
98. 98 Mutex::new(Vec::new()),
99. 99 ));
100. 100 let dir = tempfile::tempdir().unwrap();
101. 101 let cwd = dir.path().to_path_buf();
102. 102 std::fs::create_dir_all(cwd.join(".vak")).unwrap();
103. 103 std::fs::write(
104. 104 cwd.join(".vak/config.toml"),
105. 105 "[memory]\nreflection = false\n[gateway]\nchat_allowlist_open = true\n",
106. 106 )
107. 107 .unwrap();
108. 108 vak_config::paths::isolate_home_for_tests();
109. 109 let core = Core::new_with_trust(cwd.clone(), true).unwrap();
110. 110 core.set_sessions_home(dir.path().join("home"));
111. 111 core.set_permission_mode(vak_config::PermissionMode::WorkspaceWrite);
112. 112 core.set_tool_worker_exe(std::path::PathBuf::from(env!(
113. 113 "CARGO_BIN_EXE_vak-tool-worker"
114. 114 )));
115. 115 core.set_provider_instance(provider);
116. 116 std::mem::forget(dir);
117. 117
118. 118 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
119. 119 let addr = listener.local_addr().unwrap();
120. 120 tokio::spawn(async move {
121. 121 axum::serve(listener, vak_server::gateway_router(core))
122. 122 .await
123. 123 .unwrap();
124. 124 });
125. 125 let mut request = serde_json::json!({
126. 126 "surface": "telegram",
127. 127 "chat": "42",
128. 128 "sender": "7",
129. 129 "text": "update the file",
130. 130 "wait": true,
131. 131 "attachments": [{
132. 132 "mime": "application/octet-stream",
133. 133 "data": base64::engine::general_purpose::STANDARD.encode(bytes),
134. 134 "filename": name,
135. 135 "kind": "document",
136. 136 }],
137. 137 });
138. 138 if accepts_files {
139. 139 request["capabilities"] = serde_json::json!({"accepts_files": true});
140. 140 }
141. 141 let response = reqwest::Client::new()
142. 142 .post(format!("http://{addr}/gateway/inbound"))
143. 143 .json(&request)
144. 144 .send()
145. 145 .await
146. 146 .unwrap();
147. 147 assert_eq!(response.status(), 200);
148. 148 response.json().await.unwrap()
149. 149}
150. 150
151. 151#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
152. 152async fn an_edited_workbook_comes_back_to_the_chat_it_came_from() {
153. 153 let reply = round_trip(
154. 154 "budget.xlsx",
155. 155 &vak_ooxml::fixtures::xlsx(),
156. 156 serde_json::json!([{"op": "set_cells", "sheet": "Budget", "cells": {"B2": 150}}]),
157. 157 true,
158. 158 )
159. 159 .await;
160. 160 let files = reply["files"].as_array().unwrap();
161. 161 assert_eq!(files.len(), 1, "{reply}");
162. 162 assert_eq!(files[0]["name"], "budget.xlsx");
163. 163 assert!(
164. 164 files[0]["caption"]
165. 165 .as_str()
166. 166 .unwrap()
167. 167 .contains("Budget: 1 changed"),
168. 168 "{reply}"
169. 169 );
170. 170 let bytes = base64::engine::general_purpose::STANDARD
171. 171 .decode(files[0]["data"].as_str().unwrap())
172. 172 .unwrap();
173. 173 let document =
174. 174 vak_ooxml::read::read(std::io::Cursor::new(bytes), vak_ooxml::Limits::default()).unwrap();
175. 175 assert!(
176. 176 document
177. 177 .units
178. 178 .iter()
179. 179 .flat_map(|unit| unit.cells.iter())
180. 180 .any(|(address, value)| address == "B2" && value == "150"),
181. 181 "the returned workbook holds the edit"
182. 182 );
183. 183}
184. 184
185. 185#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
186. 186async fn a_labelled_draft_is_held_and_a_bridge_without_files_is_told_where_it_is() {
187. 187 let labelled = round_trip(
188. 188 "memo.docx",
189. 189 &vak_ooxml::fixtures::signed_labelled_docx(),
190. 190 serde_json::json!([{"op": "replace_paragraph_text", "anchor": "p@1", "text": "Hello again"}]),
191. 191 true,
192. 192 )
193. 193 .await;
194. 194 assert!(
195. 195 labelled["files"].as_array().unwrap().is_empty(),
196. 196 "{labelled}"
197. 197 );
198. 198 assert!(
199. 199 labelled["text"]
200. 200 .as_str()
201. 201 .unwrap()
202. 202 .contains("memo.docx carries the sensitivity label Confidential"),
203. 203 "{labelled}"
204. 204 );
205. 205
206. 206 let plain = round_trip(
207. 207 "budget.xlsx",
208. 208 &vak_ooxml::fixtures::xlsx(),
209. 209 serde_json::json!([{"op": "set_cells", "sheet": "Budget", "cells": {"B2": 150}}]),
210. 210 false,
211. 211 )
212. 212 .await;
213. 213 assert!(plain["files"].as_array().unwrap().is_empty(), "{plain}");
214. 214 assert!(
215. 215 plain["text"]
216. 216 .as_str()
217. 217 .unwrap()
218. 218 .contains("The updated budget.xlsx is ready in Vak"),
219. 219 "{plain}"
220. 220 );
221. 221}
222. 222
223. 223/// The whole Telegram path against a Bot API double: a workbook arrives as
224. 224/// a document, the bridge hands it to the gateway, and the edited workbook
225. 225/// goes back to the same chat with `sendDocument`, after the text reply.
226. 226#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
227. 227async fn the_telegram_bridge_sends_the_edited_workbook_back() {
228. 228 let workbook = vak_ooxml::fixtures::xlsx();
229. 229 let (path, digest) = inbox_path("budget.xlsx", &workbook);
230. 230 let sent: Arc<Mutex<Vec<(String, String)>>> = Arc::new(Mutex::new(Vec::new()));
231. 231 let served = Arc::new(Mutex::new(false));
232. 232 let app = {
233. 233 let sent_message = sent.clone();
234. 234 let sent_document = sent.clone();
235. 235 let file = workbook.clone();
236. 236 axum::Router::new()
237. 237 .route(
238. 238 "/botbottok/getUpdates",
239. 239 axum::routing::get(move || async move {
240. 240 let mut once = served.lock().unwrap();
241. 241 let result = if *once {
242. 242 serde_json::json!([])
243. 243 } else {
244. 244 *once = true;
245. 245 serde_json::json!([{
246. 246 "update_id": 900,
247. 247 "message": {
248. 248 "chat": {"id": 4242},
249. 249 "from": {"id": 7},
250. 250 "caption": "raise food to 150",
251. 251 "document": {"file_id": "doc1", "file_name": "budget.xlsx", "mime_type": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"}
252. 252 }
253. 253 }])
254. 254 };
255. 255 axum::Json(serde_json::json!({"ok": true, "result": result}))
256. 256 }),
257. 257 )
258. 258 .route(
259. 259 "/botbottok/getFile",
260. 260 axum::routing::get(|| async {
261. 261 axum::Json(serde_json::json!({"ok": true, "result": {"file_path": "documents/budget.xlsx"}}))
262. 262 }),
263. 263 )
264. 264 .route(
265. 265 "/file/botbottok/documents/budget.xlsx",
266. 266 axum::routing::get(move || {
267. 267 let file = file.clone();
268. 268 async move { file }
269. 269 }),
270. 270 )
271. 271 .route(
272. 272 "/botbottok/sendMessage",
273. 273 axum::routing::post(move |body: String| async move {
274. 274 sent_message.lock().unwrap().push(("message".into(), body));
275. 275 axum::Json(serde_json::json!({"ok": true}))
276. 276 }),
277. 277 )
278. 278 .route(
279. 279 "/botbottok/sendDocument",
280. 280 axum::routing::post(move |body: axum::body::Bytes| async move {
281. 281 sent_document
282. 282 .lock()
283. 283 .unwrap()
284. 284 .push(("document".into(), String::from_utf8_lossy(&body).into_owned()));
285. 285 axum::Json(serde_json::json!({"ok": true}))
286. 286 }),
287. 287 )
288. 288 .layer(axum::extract::DefaultBodyLimit::max(64 * 1024 * 1024))
289. 289 };
290. 290 let telegram = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
291. 291 let telegram_addr = telegram.local_addr().unwrap();
292. 292 tokio::spawn(async move { axum::serve(telegram, app).await.unwrap() });
293. 293
294. 294 let provider = Arc::new(Scripted(
295. 295 Mutex::new(VecDeque::from([
296. 296 message(
297. 297 ContentBlock::ToolUse {
298. 298 id: "call_edit".into(),
299. 299 name: "office_apply".into(),
300. 300 input: serde_json::json!({
301. 301 "path": path,
302. 302 "base_digest": digest,
303. 303 "ops": [{"op": "set_cells", "sheet": "Budget", "cells": {"B2": 150}}],
304. 304 }),
305. 305 },
306. 306 StopReason::ToolUse,
307. 307 ),
308. 308 message(ContentBlock::text("Raised it to 150."), StopReason::EndTurn),
309. 309 ])),
310. 310 Mutex::new(Vec::new()),
311. 311 ));
312. 312 let dir = tempfile::tempdir().unwrap();
313. 313 let cwd = dir.path().to_path_buf();
314. 314 std::fs::create_dir_all(cwd.join(".vak")).unwrap();
315. 315 std::fs::write(
316. 316 cwd.join(".vak/config.toml"),
317. 317 "[memory]\nreflection = false\n[gateway]\nchat_allowlist = [\"telegram:4242\"]\n",
318. 318 )
319. 319 .unwrap();
320. 320 vak_config::paths::isolate_home_for_tests();
321. 321 let core = Core::new_with_trust(cwd.clone(), true).unwrap();
322. 322 core.set_sessions_home(dir.path().join("home"));
323. 323 core.set_permission_mode(vak_config::PermissionMode::WorkspaceWrite);
324. 324 core.set_tool_worker_exe(std::path::PathBuf::from(env!(
325. 325 "CARGO_BIN_EXE_vak-tool-worker"
326. 326 )));
327. 327 core.set_provider_instance(provider.clone());
328. 328 std::mem::forget(dir);
329. 329 let gateway = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
330. 330 let gateway_addr = gateway.local_addr().unwrap();
331. 331 tokio::spawn(async move {
332. 332 axum::serve(gateway, vak_server::gateway_router(core))
333. 333 .await
334. 334 .unwrap();
335. 335 });
336. 336
337. 337 let bridge = vak_server::surfaces::telegram::TelegramBridge {
338. 338 token_env: String::new(),
339. 339 locks_dir: None,
340. 340 api_base: format!("http://{telegram_addr}"),
341. 341 bot_token: "bottok".into(),
342. 342 gateway_url: format!("http://{gateway_addr}"),
343. 343 gateway_token: "vk_test".into(),
344. 344 bot_id: None,
345. 345 };
346. 346 assert_eq!(bridge.tick(0).await.unwrap(), 901);
347. 347
348. 348 let sent = sent.lock().unwrap();
349. 349 assert_eq!(sent.len(), 2, "{sent:?}");
350. 350 assert_eq!(sent[0].0, "message");
351. 351 assert!(sent[0].1.contains("Raised it to 150."), "{}", sent[0].1);
352. 352 let (kind, body) = &sent[1];
353. 353 assert_eq!(kind, "document");
354. 354 assert!(
355. 355 body.contains("filename=\"budget.xlsx\""),
356. 356 "the file keeps its sent name"
357. 357 );
358. 358 assert!(
359. 359 body.contains("Updated budget.xlsx: Budget: 1 changed"),
360. 360 "{body}"
361. 361 );
362. 362 assert!(
363. 363 body.contains("name=\"chat_id\"\r\n\r\n4242"),
364. 364 "back to the same chat"
365. 365 );
366. 366 let asked = serde_json::to_string(&provider.1.lock().unwrap()[0].messages).unwrap();
367. 367 assert!(
368. 368 asked.contains("raise food to 150"),
369. 369 "the file's caption is the request"
370. 370 );
371. 371}
372. 372