1. 1#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
2. 2
3. 3use std::collections::{HashMap, VecDeque};
4. 4use std::sync::{Arc, Mutex};
5. 5
6. 6use tokio::sync::mpsc;
7. 7use tokio_util::sync::CancellationToken;
8. 8
9. 9use vak_agent::AutoApprove;
10. 10use vak_flow::{Executor, ExecutorDeps, FlowOutcome, FlowState};
11. 11use vak_llm::stream;
12. 12use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, StopReason, Usage};
13. 13use vak_llm::{EventStream, LlmError, Provider};
14. 14use vak_permission::{Mode, PermissionEngine};
15. 15use vak_tools::bash::BashTool;
16. 16
17. 17/// Serves a fixed sequence per routing key (the first user text).
18. 18struct TaggedScripted {
19. 19 routes: Mutex<HashMap<String, VecDeque<AssistantMessage>>>,
20. 20}
21. 21
22. 22impl TaggedScripted {
23. 23 fn route_for(request: &ChatRequest) -> String {
24. 24 request
25. 25 .messages
26. 26 .iter()
27. 27 .find(|m| m.role == vak_llm::types::Role::User)
28. 28 .map(|m| m.text_content())
29. 29 .unwrap_or_default()
30. 30 }
31. 31}
32. 32
33. 33#[async_trait::async_trait]
34. 34impl Provider for TaggedScripted {
35. 35 fn name(&self) -> &str {
36. 36 "scripted"
37. 37 }
38. 38
39. 39 async fn stream(
40. 40 &self,
41. 41 request: ChatRequest,
42. 42 _cancel: CancellationToken,
43. 43 ) -> Result<EventStream, LlmError> {
44. 44 let key = Self::route_for(&request);
45. 45 let next = self
46. 46 .routes
47. 47 .lock()
48. 48 .unwrap()
49. 49 .get_mut(&key)
50. 50 .and_then(|d| d.pop_front());
51. 51 let (mut sink, rx) = stream::channel(64);
52. 52 match next {
53. 53 Some(m) => {
54. 54 sink.push(stream::StreamEvent::Start { partial: m.clone() });
55. 55 sink.close_message(m).await;
56. 56 }
57. 57 None => {
58. 58 sink.close_error(LlmError::Parse(format!("exhausted: {key}")))
59. 59 .await
60. 60 }
61. 61 }
62. 62 Ok(rx)
63. 63 }
64. 64}
65. 65
66. 66fn text(t: &str) -> AssistantMessage {
67. 67 AssistantMessage {
68. 68 content: vec![ContentBlock::text(t)],
69. 69 stop_reason: StopReason::EndTurn,
70. 70 usage: Usage::default(),
71. 71 model: "test-model".into(),
72. 72 response_id: None,
73. 73 }
74. 74}
75. 75
76. 76fn make_executor(provider: Arc<TaggedScripted>, state_path: std::path::PathBuf) -> Executor {
77. 77 // Bash records workspace file changes. Sharing the system temp directory
78. 78 // lets concurrent tests' files leak into the output used by prompt substitution.
79. 79 let workspace = tempfile::tempdir().unwrap().keep();
80. 80 make_executor_with_policy(
81. 81 provider,
82. 82 state_path,
83. 83 Mode::FullAccess,
84. 84 Some(Arc::new(AutoApprove)),
85. 85 workspace,
86. 86 )
87. 87}
88. 88
89. 89fn make_executor_with_policy(
90. 90 provider: Arc<TaggedScripted>,
91. 91 state_path: std::path::PathBuf,
92. 92 mode: Mode,
93. 93 approver: Option<Arc<dyn vak_agent::Approver>>,
94. 94 cwd: std::path::PathBuf,
95. 95) -> Executor {
96. 96 make_executor_with_outcome(provider, state_path, mode, approver, cwd, None)
97. 97}
98. 98
99. 99fn make_executor_with_outcome(
100. 100 provider: Arc<TaggedScripted>,
101. 101 state_path: std::path::PathBuf,
102. 102 mode: Mode,
103. 103 approver: Option<Arc<dyn vak_agent::Approver>>,
104. 104 cwd: std::path::PathBuf,
105. 105 outcome: Option<vak_intent::OutcomeSpec>,
106. 106) -> Executor {
107. 107 let dir = tempfile::tempdir().unwrap();
108. 108 let home = dir.path().join("home");
109. 109 std::fs::create_dir_all(&home).unwrap();
110. 110 std::mem::forget(dir);
111. 111 Executor::new(ExecutorDeps {
112. 112 prompt_layers: Vec::new(),
113. 113 provider,
114. 114 system_prompt: "sys".into(),
115. 115 model: "test-model".into(),
116. 116 tools: vec![Arc::new(BashTool)],
117. 117 read_only_tools: vec![],
118. 118 max_turns: 4,
119. 119 outcome,
120. 120 max_retries: 0,
121. 121 retry_base_backoff_ms: 100,
122. 122 request_timeout: Some(std::time::Duration::from_secs(600)),
123. 123 circuit_breaker: None,
124. 124 run_retry_attempts: 0,
125. 125 run_retry_base_backoff_ms: 1000,
126. 126 dispatch_ceiling: 1,
127. 127 spend_gate: None,
128. 128 permission: Some(Arc::new(PermissionEngine::default())),
129. 129 mode,
130. 130 approval_mode: vak_agent::ApprovalMode::Ask,
131. 131 approver,
132. 132 sandbox: None,
133. 133 cwd,
134. 134 sessions_home: home,
135. 135 parent_session_id: "flow-parent".into(),
136. 136 state_path,
137. 137 agent_identity: None,
138. 138 conversation_context: None,
139. 139 work: None,
140. 140 })
141. 141}
142. 142
143. 143#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
144. 144async fn admitted_outcome_is_persisted_into_flow_state() {
145. 145 let workspace = tempfile::tempdir().unwrap();
146. 146 let toml = r#"
147. 147[flow]
148. 148name = "outcome"
149. 149
150. 150[[nodes]]
151. 151id = "report"
152. 152type = "bash"
153. 153command = "echo admitted"
154. 154"#;
155. 155 let flow = vak_flow::parse_flow(toml).unwrap();
156. 156 let expected = vak_intent::OutcomeSpec {
157. 157 schema_version: 1,
158. 158 revision: 4,
159. 159 objective: "preserve the admitted objective".into(),
160. 160 assumptions: vec!["the workspace is available".into()],
161. 161 requirements: Vec::new(),
162. 162 resolver_version: 1,
163. 163 evidence_max_age_secs: Some(3600),
164. 164 acts: Default::default(),
165. 165 stop: Default::default(),
166. 166 max_turns: Some(2),
167. 167 };
168. 168 let provider = Arc::new(TaggedScripted {
169. 169 routes: Mutex::new(HashMap::new()),
170. 170 });
171. 171 let state_path = workspace.path().join("state.json");
172. 172 let mut state = FlowState {
173. 173 run_id: "outcome-run".into(),
174. 174 flow_name: "outcome".into(),
175. 175 definition_toml: toml.into(),
176. 176 started_at: chrono::Utc::now(),
177. 177 outcome: None,
178. 178 nodes: Default::default(),
179. 179 };
180. 180 let executor = make_executor_with_outcome(
181. 181 provider,
182. 182 state_path.clone(),
183. 183 Mode::FullAccess,
184. 184 Some(Arc::new(AutoApprove)),
185. 185 workspace.path().to_path_buf(),
186. 186 Some(expected.clone()),
187. 187 );
188. 188
189. 189 let _ = drain_run(&executor, &flow, &mut state).await;
190. 190 assert_eq!(state.outcome, Some(expected));
191. 191 let persisted = std::fs::read_to_string(state_path).unwrap();
192. 192 assert!(persisted.contains("preserve the admitted objective"));
193. 193}
194. 194
195. 195#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
196. 196async fn bash_node_obeys_permission_decision() {
197. 197 let workspace = tempfile::tempdir().unwrap();
198. 198 let toml = r#"
199. 199[flow]
200. 200name = "permission-gate"
201. 201
202. 202[[nodes]]
203. 203id = "blocked"
204. 204type = "bash"
205. 205command = "echo escaped > should-not-exist.txt"
206. 206"#;
207. 207 let flow = vak_flow::parse_flow(toml).unwrap();
208. 208 let provider = Arc::new(TaggedScripted {
209. 209 routes: Mutex::new(HashMap::new()),
210. 210 });
211. 211 let state_path = workspace.path().join("state.json");
212. 212 let mut state = FlowState {
213. 213 run_id: "permission-run".into(),
214. 214 flow_name: "permission-gate".into(),
215. 215 definition_toml: toml.into(),
216. 216 started_at: chrono::Utc::now(),
217. 217 outcome: None,
218. 218 nodes: Default::default(),
219. 219 };
220. 220 let executor = make_executor_with_policy(
221. 221 provider,
222. 222 state_path,
223. 223 Mode::WorkspaceWrite,
224. 224 None,
225. 225 workspace.path().to_path_buf(),
226. 226 );
227. 227
228. 228 let outcome = drain_run(&executor, &flow, &mut state).await;
229. 229 match outcome {
230. 230 FlowOutcome::Failed { node, reason, .. } => {
231. 231 assert_eq!(node, "blocked");
232. 232 assert!(reason.contains("no approver available"), "{reason}");
233. 233 }
234. 234 other => panic!("expected permission failure, got {other:?}"),
235. 235 }
236. 236 assert!(!workspace.path().join("should-not-exist.txt").exists());
237. 237}
238. 238
239. 239async fn drain_run(
240. 240 executor: &Executor,
241. 241 flow: &vak_flow::FlowDef,
242. 242 state: &mut FlowState,
243. 243) -> FlowOutcome {
244. 244 let (tx, mut rx) = mpsc::channel(256);
245. 245 let drainer = tokio::spawn(async move { while rx.recv().await.is_some() {} });
246. 246 let outcome = executor
247. 247 .run(flow, state, CancellationToken::new(), tx)
248. 248 .await;
249. 249 drainer.await.unwrap();
250. 250 outcome
251. 251}
252. 252
253. 253#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
254. 254async fn chain_runs_with_substitution_and_merge() {
255. 255 let toml = r#"
256. 256[flow]
257. 257name = "chain"
258. 258
259. 259[[nodes]]
260. 260id = "greet"
261. 261type = "bash"
262. 262command = "echo hello"
263. 263
264. 264[[nodes]]
265. 265id = "shout"
266. 266type = "agent"
267. 267prompt = "SHOUT {{greet}}"
268. 268
269. 269[[nodes]]
270. 270id = "done"
271. 271type = "merge"
272. 272deps = ["shout"]
273. 273"#;
274. 274 let flow = vak_flow::parse_flow(toml).unwrap();
275. 275
276. 276 let mut routes = HashMap::new();
277. 277 routes.insert(
278. 278 "SHOUT [stdout]\nhello\n\n".to_string(),
279. 279 VecDeque::from(vec![text("HELLO")]),
280. 280 );
281. 281 let provider = Arc::new(TaggedScripted {
282. 282 routes: Mutex::new(routes),
283. 283 });
284. 284
285. 285 let state_path = std::env::temp_dir().join(format!("vak-flow-{}.json", uuid_like()));
286. 286 let mut state = FlowState {
287. 287 run_id: "r1".into(),
288. 288 flow_name: "chain".into(),
289. 289 definition_toml: toml.into(),
290. 290 started_at: chrono::Utc::now(),
291. 291 outcome: None,
292. 292 nodes: Default::default(),
293. 293 };
294. 294
295. 295 let executor = make_executor(provider, state_path.clone());
296. 296 let outcome = drain_run(&executor, &flow, &mut state).await;
297. 297
298. 298 match outcome {
299. 299 FlowOutcome::Completed { outputs } => {
300. 300 assert!(outputs.get("greet").is_some_and(|o| o.contains("hello")));
301. 301 assert_eq!(outputs.get("shout").map(String::as_str), Some("HELLO"));
302. 302 let merge = outputs.get("done").expect("merge output");
303. 303 assert!(merge.contains("[shout] ok"));
304. 304 assert!(merge.contains("HELLO"));
305. 305 }
306. 306 other => panic!("expected completed, got {other:?}"),
307. 307 }
308. 308 assert_eq!(
309. 309 state.nodes.get("greet").unwrap().status,
310. 310 vak_flow::NodeStatus::Completed
311. 311 );
312. 312 let _ = std::fs::remove_file(&state_path);
313. 313}
314. 314
315. 315#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
316. 316async fn required_failure_fails_flow_and_skips_downstream() {
317. 317 let toml = r#"
318. 318[flow]
319. 319name = "strict"
320. 320
321. 321[[nodes]]
322. 322id = "boom"
323. 323type = "bash"
324. 324command = "exit 3"
325. 325required = true
326. 326
327. 327[[nodes]]
328. 328id = "after"
329. 329type = "agent"
330. 330prompt = "never runs"
331. 331deps = ["boom"]
332. 332"#;
333. 333 let flow = vak_flow::parse_flow(toml).unwrap();
334. 334 let provider = Arc::new(TaggedScripted {
335. 335 routes: Mutex::new(HashMap::new()),
336. 336 });
337. 337 let state_path = std::env::temp_dir().join(format!("vak-flow-{}.json", uuid_like()));
338. 338 let mut state = FlowState {
339. 339 run_id: "r2".into(),
340. 340 flow_name: "strict".into(),
341. 341 definition_toml: toml.into(),
342. 342 started_at: chrono::Utc::now(),
343. 343 outcome: None,
344. 344 nodes: Default::default(),
345. 345 };
346. 346 let executor = make_executor(provider, state_path.clone());
347. 347 let outcome = drain_run(&executor, &flow, &mut state).await;
348. 348
349. 349 match outcome {
350. 350 FlowOutcome::Failed { node, reason, .. } => {
351. 351 assert_eq!(node, "boom");
352. 352 assert!(reason.contains("exit code: 3"));
353. 353 }
354. 354 other => panic!("expected failed, got {other:?}"),
355. 355 }
356. 356 assert_eq!(
357. 357 state.nodes.get("after").unwrap().status,
358. 358 vak_flow::NodeStatus::Skipped
359. 359 );
360. 360 let _ = std::fs::remove_file(&state_path);
361. 361}
362. 362
363. 363#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
364. 364async fn optional_failure_lets_merge_report_partial_outcome() {
365. 365 let toml = r#"
366. 366[flow]
367. 367name = "lenient"
368. 368
369. 369[[nodes]]
370. 370id = "flaky"
371. 371type = "bash"
372. 372command = "exit 9"
373. 373required = false
374. 374
375. 375[[nodes]]
376. 376id = "report"
377. 377type = "merge"
378. 378deps = ["flaky"]
379. 379"#;
380. 380 let flow = vak_flow::parse_flow(toml).unwrap();
381. 381 let provider = Arc::new(TaggedScripted {
382. 382 routes: Mutex::new(HashMap::new()),
383. 383 });
384. 384 let state_path = std::env::temp_dir().join(format!("vak-flow-{}.json", uuid_like()));
385. 385 let mut state = FlowState {
386. 386 run_id: "r3".into(),
387. 387 flow_name: "lenient".into(),
388. 388 definition_toml: toml.into(),
389. 389 started_at: chrono::Utc::now(),
390. 390 outcome: None,
391. 391 nodes: Default::default(),
392. 392 };
393. 393 let executor = make_executor(provider, state_path.clone());
394. 394 let outcome = drain_run(&executor, &flow, &mut state).await;
395. 395
396. 396 match outcome {
397. 397 FlowOutcome::Completed { outputs } => {
398. 398 let report = outputs.get("report").expect("merge ran despite failure");
399. 399 assert!(report.contains("[flaky] unavailable"));
400. 400 }
401. 401 other => panic!("expected completed with partial report, got {other:?}"),
402. 402 }
403. 403 assert_eq!(
404. 404 state.nodes.get("flaky").unwrap().status,
405. 405 vak_flow::NodeStatus::Failed
406. 406 );
407. 407 let _ = std::fs::remove_file(&state_path);
408. 408}
409. 409
410. 410#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
411. 411async fn resume_skips_completed_nodes() {
412. 412 let toml = r#"
413. 413[flow]
414. 414name = "resumable"
415. 415
416. 416[[nodes]]
417. 417id = "first"
418. 418type = "bash"
419. 419command = "echo one"
420. 420
421. 421[[nodes]]
422. 422id = "second"
423. 423type = "bash"
424. 424command = "echo two"
425. 425deps = ["first"]
426. 426"#;
427. 427 let flow = vak_flow::parse_flow(toml).unwrap();
428. 428 let provider = Arc::new(TaggedScripted {
429. 429 routes: Mutex::new(HashMap::new()),
430. 430 });
431. 431 let state_path = std::env::temp_dir().join(format!("vak-flow-{}.json", uuid_like()));
432. 432
433. 433 // Pre-seed state as if "first" already completed in an earlier run.
434. 434 let mut state = FlowState {
435. 435 run_id: "r4".into(),
436. 436 flow_name: "resumable".into(),
437. 437 definition_toml: toml.into(),
438. 438 started_at: chrono::Utc::now(),
439. 439 outcome: None,
440. 440 nodes: [(
441. 441 "first".to_string(),
442. 442 vak_flow::NodeResult {
443. 443 status: vak_flow::NodeStatus::Completed,
444. 444 output: "one\n".into(),
445. 445 },
446. 446 )]
447. 447 .into_iter()
448. 448 .collect(),
449. 449 };
450. 450 std::fs::write(&state_path, serde_json::to_string_pretty(&state).unwrap()).unwrap();
451. 451
452. 452 let executor = make_executor(provider.clone(), state_path.clone());
453. 453 let outcome = drain_run(&executor, &flow, &mut state).await;
454. 454 assert!(matches!(outcome, FlowOutcome::Completed { .. }));
455. 455
456. 456 // "second" must have executed; "first" must not have re-executed.
457. 457 assert_eq!(
458. 458 state.nodes.get("second").unwrap().status,
459. 459 vak_flow::NodeStatus::Completed
460. 460 );
461. 461 assert!(state.nodes.get("second").unwrap().output.contains("two"));
462. 462 assert_eq!(state.nodes.get("first").unwrap().output.trim(), "one");
463. 463
464. 464 // The persisted ledger reflects both.
465. 465 let persisted: FlowState =
466. 466 serde_json::from_str(&std::fs::read_to_string(&state_path).unwrap()).unwrap();
467. 467 assert_eq!(persisted.nodes.len(), 2);
468. 468 let _ = std::fs::remove_file(&state_path);
469. 469}
470. 470
471. 471fn uuid_like() -> String {
472. 472 use std::sync::atomic::{AtomicU32, Ordering};
473. 473 static C: AtomicU32 = AtomicU32::new(0);
474. 474 format!("u{}", C.fetch_add(1, Ordering::Relaxed))
475. 475}
476. 476
477. 477#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
478. 478async fn done_contract_failure_rejects_node() {
479. 479 let workspace = tempfile::tempdir().unwrap();
480. 480 let toml = r#"
481. 481[flow]
482. 482name = "contract"
483. 483
484. 484[[nodes]]
485. 485id = "build"
486. 486type = "bash"
487. 487command = "echo built > built.txt"
488. 488accept = ["verify: test -f built.txt", "verify: test -f missing-artifact.txt"]
489. 489
490. 490[[nodes]]
491. 491id = "after"
492. 492type = "merge"
493. 493deps = ["build"]
494. 494"#;
495. 495 let flow = vak_flow::parse_flow(toml).unwrap();
496. 496 let provider = Arc::new(TaggedScripted {
497. 497 routes: Mutex::new(HashMap::new()),
498. 498 });
499. 499 let state_path = workspace.path().join("state.json");
500. 500 let executor = make_executor_with_policy(
501. 501 provider,
502. 502 state_path.clone(),
503. 503 Mode::FullAccess,
504. 504 None,
505. 505 workspace.path().to_path_buf(),
506. 506 );
507. 507
508. 508 let (tx, mut rx) = mpsc::channel(256);
509. 509 tokio::spawn(async move { while rx.recv().await.is_some() {} });
510. 510 let mut state = FlowState {
511. 511 run_id: "r".into(),
512. 512 flow_name: "contract".into(),
513. 513 definition_toml: toml.into(),
514. 514 started_at: chrono::Utc::now(),
515. 515 outcome: None,
516. 516 nodes: Default::default(),
517. 517 };
518. 518 let outcome = executor
519. 519 .run(&flow, &mut state, CancellationToken::new(), tx)
520. 520 .await;
521. 521 match outcome {
522. 522 FlowOutcome::Failed { node, reason, .. } => {
523. 523 assert_eq!(node, "build");
524. 524 assert!(reason.contains("done-contract failed"), "{reason}");
525. 525 assert!(reason.contains("missing-artifact.txt"), "{reason}");
526. 526 }
527. 527 other => panic!("expected contract failure, got {other:?}"),
528. 528 }
529. 529
530. 530 // Ledger marks the node Failed despite the command itself succeeding.
531. 531 let st: FlowState =
532. 532 serde_json::from_str(&std::fs::read_to_string(&state_path).unwrap()).unwrap();
533. 533 assert_eq!(
534. 534 st.nodes.get("build").map(|r| r.status),
535. 535 Some(vak_flow::NodeStatus::Failed)
536. 536 );
537. 537}
538. 538
539. 539#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
540. 540async fn done_contract_pass_keeps_node_green() {
541. 541 let workspace = tempfile::tempdir().unwrap();
542. 542 let toml = r#"
543. 543[flow]
544. 544name = "contract-ok"
545. 545
546. 546[[nodes]]
547. 547id = "make"
548. 548type = "bash"
549. 549command = "echo v1 > artifact.txt"
550. 550accept = ["verify: grep -q v1 artifact.txt"]
551. 551"#;
552. 552 let flow = vak_flow::parse_flow(toml).unwrap();
553. 553 let provider = Arc::new(TaggedScripted {
554. 554 routes: Mutex::new(HashMap::new()),
555. 555 });
556. 556 let state_path = workspace.path().join("state.json");
557. 557 let executor = make_executor_with_policy(
558. 558 provider,
559. 559 state_path.clone(),
560. 560 Mode::FullAccess,
561. 561 None,
562. 562 workspace.path().to_path_buf(),
563. 563 );
564. 564 let (tx, mut rx) = mpsc::channel(256);
565. 565 tokio::spawn(async move { while rx.recv().await.is_some() {} });
566. 566 let mut state = FlowState {
567. 567 run_id: "r".into(),
568. 568 flow_name: "contract-ok".into(),
569. 569 definition_toml: toml.into(),
570. 570 started_at: chrono::Utc::now(),
571. 571 outcome: None,
572. 572 nodes: Default::default(),
573. 573 };
574. 574 let outcome = executor
575. 575 .run(&flow, &mut state, CancellationToken::new(), tx)
576. 576 .await;
577. 577 assert!(matches!(outcome, FlowOutcome::Completed { .. }));
578. 578 let st: FlowState =
579. 579 serde_json::from_str(&std::fs::read_to_string(&state_path).unwrap()).unwrap();
580. 580 assert_eq!(
581. 581 st.nodes.get("make").map(|r| r.status),
582. 582 Some(vak_flow::NodeStatus::Completed)
583. 583 );
584. 584}
585. 585