Vakyartha documentation

1#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
2
3use std::collections::VecDeque;
4use std::sync::{Arc, Mutex};
5
6use tokio::sync::mpsc;
7use tokio_util::sync::CancellationToken;
8
9use tempfile::tempdir;
10
11use vak_agent::AutoApprove;
12use vak_flow::{
13 ExecutorDeps, PLANNER_SYSTEM, PlanOutcome, ToolCatalogEntry, build_planner_prompt,
14 extract_toml, plan_and_run, sanitize_basic_string_newlines,
15};
16use vak_llm::stream;
17use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, StopReason, Usage};
18use vak_llm::{EventStream, LlmError, Provider};
19use vak_permission::{Mode, PermissionEngine};
20use vak_tools::bash::BashTool;
21
22#[test]
23fn planner_prompt_includes_task_and_catalog() {
24 let catalog = vec![
25 ToolCatalogEntry {
26 name: "bash".into(),
27 description: "Execute a shell command".into(),
28 },
29 ToolCatalogEntry {
30 name: "task".into(),
31 description: "Delegate to a worker".into(),
32 },
33 ];
34 let prompt = build_planner_prompt("migrate the config module", &catalog);
35 assert!(prompt.contains("migrate the config module"));
36 assert!(prompt.contains("- bash: Execute a shell command"));
37 assert!(prompt.contains("- task: Delegate to a worker"));
38}
39
40#[test]
41fn extract_toml_handles_fenced_raw_and_garbage() {
42 let fenced = "Here is the plan:\n```toml\n[flow]\nname = \"x\"\n[[nodes]]\nid=\"a\"\ntype=\"bash\"\ncommand=\"echo hi\"\n```\ndone";
43 let extracted = extract_toml(fenced).unwrap();
44 assert!(extracted.contains("[flow]"));
45 assert!(!extracted.contains("```"));
46
47 let raw = "[flow]\nname = \"y\"\n\n[[nodes]]\nid = \"a\"\ntype = \"bash\"\ncommand = \"echo\"";
48 assert_eq!(extract_toml(raw).as_deref(), Some(raw));
49
50 assert!(extract_toml("I cannot plan this task, sorry.").is_none());
51}
52
53struct ScriptedPlanner {
54 /// Responses consumed in order across ALL requests (planner + children).
55 responses: Mutex<VecDeque<ScriptedResponse>>,
56}
57
58enum ScriptedResponse {
59 Text(String),
60 Error(LlmError),
61}
62
63#[async_trait::async_trait]
64impl Provider for ScriptedPlanner {
65 fn name(&self) -> &str {
66 "scripted"
67 }
68
69 async fn stream(
70 &self,
71 _request: ChatRequest,
72 _cancel: CancellationToken,
73 ) -> Result<EventStream, LlmError> {
74 let next = self.responses.lock().unwrap().pop_front();
75 let (mut sink, rx) = stream::channel(64);
76 match next {
77 Some(ScriptedResponse::Text(t)) => {
78 let msg = AssistantMessage {
79 content: vec![ContentBlock::text(t.clone())],
80 stop_reason: StopReason::EndTurn,
81 usage: Usage::default(),
82 model: "test-model".into(),
83 response_id: None,
84 };
85 sink.push(stream::StreamEvent::Start {
86 partial: msg.clone(),
87 });
88 sink.close_message(msg).await;
89 }
90 Some(ScriptedResponse::Error(e)) => sink.close_error(e).await,
91 None => sink.close_error(LlmError::Parse("exhausted".into())).await,
92 }
93 Ok(rx)
94 }
95}
96
97const GOOD_PLAN: &str = "\
98[flow]
99name = \"planned\"
100
101[[nodes]]
102id = \"probe\"
103type = \"bash\"
104command = \"echo planned-ok\"
105
106[[nodes]]
107id = \"report\"
108type = \"merge\"
109deps = [\"probe\"]
110";
111
112fn fenced(plan: &str) -> String {
113 format!("Sure, here is the plan:\n```toml\n{plan}\n```\n")
114}
115
116fn make_deps(provider: Arc<ScriptedPlanner>) -> Arc<ExecutorDeps> {
117 let dir = tempdir().unwrap();
118 let home = dir.path().join("home");
119 std::fs::create_dir_all(&home).unwrap();
120 std::mem::forget(dir);
121 Arc::new(ExecutorDeps {
122 prompt_layers: Vec::new(),
123 provider,
124 system_prompt: "sys".into(),
125 model: "test-model".into(),
126 tools: vec![Arc::new(BashTool)],
127 read_only_tools: vec![],
128 max_turns: 4,
129 outcome: None,
130 max_retries: 0,
131 retry_base_backoff_ms: 100,
132 request_timeout: Some(std::time::Duration::from_secs(600)),
133 circuit_breaker: None,
134 run_retry_attempts: 0,
135 run_retry_base_backoff_ms: 1000,
136 dispatch_ceiling: 1,
137 spend_gate: None,
138 permission: Some(Arc::new(PermissionEngine::default())),
139 mode: Mode::FullAccess,
140 approval_mode: vak_agent::ApprovalMode::Ask,
141 approver: Some(Arc::new(AutoApprove)),
142 sandbox: None,
143 cwd: std::env::temp_dir(),
144 sessions_home: home.clone(),
145 parent_session_id: "plan-parent".into(),
146 state_path: home.join("flow-runs/plan"),
147 agent_identity: None,
148 conversation_context: None,
149 work: None,
150 })
151}
152
153async fn drain(f: impl std::future::Future<Output = PlanOutcome>) -> PlanOutcome {
154 let (tx, mut rx) = mpsc::channel::<String>(256);
155 let d = tokio::spawn(async move { while rx.recv().await.is_some() {} });
156 drop(tx);
157 let out = f.await;
158 d.await.unwrap();
159 out
160}
161
162#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
163async fn valid_plan_executes_to_completion() {
164 let provider = Arc::new(ScriptedPlanner {
165 responses: Mutex::new(VecDeque::from(vec![ScriptedResponse::Text(fenced(
166 GOOD_PLAN,
167 ))])),
168 });
169 let deps = make_deps(provider);
170
171 let outcome = drain(plan_and_run(
172 deps,
173 "run the planned probe",
174 CancellationToken::new(),
175 mpsc::channel(64).0,
176 ))
177 .await;
178
179 match outcome {
180 PlanOutcome::Completed { outputs, attempts } => {
181 assert_eq!(attempts, 1);
182 assert!(
183 outputs
184 .get("probe")
185 .is_some_and(|o| o.contains("planned-ok"))
186 );
187 }
188 other => panic!("expected completed, got {other:?}"),
189 }
190}
191
192#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
193async fn garbage_plan_fails_closed_without_execution() {
194 let provider = Arc::new(ScriptedPlanner {
195 responses: Mutex::new(VecDeque::from(vec![ScriptedResponse::Text(
196 "I'm sorry, I cannot produce a TOML plan for that.".into(),
197 )])),
198 });
199 let deps = make_deps(provider);
200
201 let outcome = drain(plan_and_run(
202 deps,
203 "impossible task",
204 CancellationToken::new(),
205 mpsc::channel(64).0,
206 ))
207 .await;
208
209 match outcome {
210 PlanOutcome::PlanningFailed { reason } => {
211 assert!(reason.contains("no TOML plan"));
212 }
213 other => panic!("expected planning_failed, got {other:?}"),
214 }
215}
216
217#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
218async fn structurally_invalid_plan_fails_closed() {
219 // A cycle must be rejected by validation, never executed.
220 let cyclic = "\
221[flow]
222name = \"cyclic\"
223
224[[nodes]]
225id = \"a\"
226type = \"bash\"
227command = \"echo a\"
228deps = [\"b\"]
229
230[[nodes]]
231id = \"b\"
232type = \"bash\"
233command = \"echo b\"
234deps = [\"a\"]
235";
236 let provider = Arc::new(ScriptedPlanner {
237 responses: Mutex::new(VecDeque::from([ScriptedResponse::Text(fenced(cyclic))])),
238 });
239 let deps = make_deps(provider);
240
241 let outcome = drain(plan_and_run(
242 deps,
243 "circular task",
244 CancellationToken::new(),
245 mpsc::channel(64).0,
246 ))
247 .await;
248
249 match outcome {
250 PlanOutcome::PlanningFailed { reason } => {
251 assert!(reason.contains("cycle"), "got: {reason}");
252 }
253 other => panic!("expected planning_failed, got {other:?}"),
254 }
255}
256
257#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
258async fn failed_execution_triggers_exactly_one_replan() {
259 // Attempt 1: plan with a required node that fails at runtime.
260 // Attempt 2 (replan): clean plan that succeeds.
261 let failing_plan = "\
262[flow]
263name = \"attempt-one\"
264
265[[nodes]]
266id = \"boom\"
267type = \"bash\"
268command = \"exit 5\"
269required = true
270
271[[nodes]]
272id = \"report\"
273type = \"merge\"
274deps = [\"boom\"]
275";
276 let provider = Arc::new(ScriptedPlanner {
277 responses: Mutex::new(VecDeque::from(vec![
278 ScriptedResponse::Text(fenced(failing_plan)),
279 ScriptedResponse::Text(fenced(GOOD_PLAN)),
280 ])),
281 });
282 let deps = make_deps(provider);
283
284 let outcome = drain(plan_and_run(
285 deps,
286 "flaky task",
287 CancellationToken::new(),
288 mpsc::channel(256).0,
289 ))
290 .await;
291
292 match outcome {
293 PlanOutcome::Completed { attempts, .. } => {
294 assert_eq!(attempts, 2, "exactly one replan allowed");
295 }
296 other => panic!("expected completed after replan, got {other:?}"),
297 }
298}
299
300#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
301async fn replan_budget_is_bounded_at_one_retry() {
302 let bad_plan = "\
303[flow]
304name = \"always-fails\"
305
306[[nodes]]
307id = \"boom\"
308type = \"bash\"
309command = \"exit 7\"
310required = true
311";
312 let provider = Arc::new(ScriptedPlanner {
313 responses: Mutex::new(VecDeque::from(vec![
314 ScriptedResponse::Text(fenced(bad_plan)),
315 ScriptedResponse::Text(fenced(bad_plan)),
316 // A third plan would be accepted here if the budget were unbounded.
317 ScriptedResponse::Text(fenced(GOOD_PLAN)),
318 ])),
319 });
320 let deps = make_deps(provider);
321
322 let outcome = drain(plan_and_run(
323 deps,
324 "doomed task",
325 CancellationToken::new(),
326 mpsc::channel(256).0,
327 ))
328 .await;
329
330 match outcome {
331 PlanOutcome::Failed { node, .. } => {
332 assert_eq!(node, "boom");
333 }
334 other => panic!("expected failed after budget exhausted, got {other:?}"),
335 }
336}
337
338#[test]
339fn planner_system_prompt_is_compact() {
340 // Prompt discipline: the planner system prompt stays small.
341 let tokens_estimate = PLANNER_SYSTEM.len() / 4;
342 assert!(
343 tokens_estimate < 500,
344 "planner system prompt too large (~{tokens_estimate} tokens)"
345 );
346}
347
348#[test]
349fn sanitizer_escapes_newlines_in_basic_strings_only() {
350 let doc = "[flow]\nname = \"x\"\n\n[[nodes]]\nid = \"a\"\ntype = \"bash\"\ncommand = \"echo one\necho two\"\n";
351 let fixed = sanitize_basic_string_newlines(doc);
352 assert_eq!(
353 fixed,
354 "[flow]\nname = \"x\"\n\n[[nodes]]\nid = \"a\"\ntype = \"bash\"\ncommand = \"echo one\\necho two\"\n"
355 );
356 assert!(vak_flow::parse_flow(&fixed).is_ok());
357
358 // Valid documents pass through byte-for-byte.
359 assert_eq!(sanitize_basic_string_newlines(GOOD_PLAN), GOOD_PLAN);
360
361 // Multi-line basic strings are untouched (already valid TOML).
362 let multi = "prompt = \"\"\"a\nb\"\"\"\n";
363 assert_eq!(sanitize_basic_string_newlines(multi), multi);
364
365 // Literal strings and comments are untouched.
366 let lit = "command = 'echo hi'\n# comment with \" quote\nname = \"z\"\n";
367 assert_eq!(sanitize_basic_string_newlines(lit), lit);
368
369 // Escaped quotes/backslashes inside basic strings survive.
370 let escaped = "prompt = \"said \\\"hi\\\" then \\\\ broke\ninto two\"\n";
371 assert_eq!(
372 sanitize_basic_string_newlines(escaped),
373 "prompt = \"said \\\"hi\\\" then \\\\ broke\\ninto two\"\n"
374 );
375}
376
377#[test]
378fn sanitizer_escapes_nested_quotes_but_keeps_closers() {
379 // Nested shell quotes are escaped; real closers survive.
380 let doc = "[flow]\nname = \"q\"\n\n[[nodes]]\nid = \"a\"\ntype = \"bash\"\ncommand = \"test -z \"$(grep x f)\" && echo \"done\"\"\n";
381 let fixed = sanitize_basic_string_newlines(doc);
382 assert!(
383 vak_flow::parse_flow(&fixed).is_ok(),
384 "sanitized doc must parse: {fixed}"
385 );
386 assert!(fixed.contains("\\\"$(grep x f)\\\""));
387 assert!(fixed.contains("\\\"done\\\""));
388
389 // Closing quotes followed by comment / whitespace+comma stay closers.
390 let ok = "name = \"x\" # trailing\nnodes-note = [ \"a\" , \"b\" ]\n";
391 assert_eq!(sanitize_basic_string_newlines(ok), ok);
392}
393
394#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
395async fn multiline_basic_string_plan_is_sanitized_and_executes() {
396 // Regression: live planner runs emitted raw newlines inside
397 // `prompt = "..."`, which is invalid TOML. The structural sanitizer must
398 // repair it so the plan parses and executes instead of failing closed.
399 let plan = "[flow]\nname = \"multiline\"\n\n[[nodes]]\nid = \"probe\"\ntype = \"bash\"\ncommand = \"echo line-one\necho line-two\"\n\n[[nodes]]\nid = \"report\"\ntype = \"merge\"\ndeps = [\"probe\"]\n";
400 let provider = Arc::new(ScriptedPlanner {
401 responses: Mutex::new(VecDeque::from([ScriptedResponse::Text(fenced(plan))])),
402 });
403 let deps = make_deps(provider);
404
405 let outcome = drain(plan_and_run(
406 deps,
407 "multiline task",
408 CancellationToken::new(),
409 mpsc::channel(64).0,
410 ))
411 .await;
412
413 match outcome {
414 PlanOutcome::Completed { outputs, .. } => {
415 let out = outputs.get("probe").unwrap();
416 assert!(out.contains("line-one") && out.contains("line-two"));
417 }
418 other => panic!("expected completed after sanitize, got {other:?}"),
419 }
420}
421
422#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
423async fn transient_planner_failure_is_retried() {
424 // Invariant 7: the planner bypasses the loop's retry machinery, so it
425 // must retry transient provider failures itself instead of failing the
426 // whole run closed.
427 let provider = Arc::new(ScriptedPlanner {
428 responses: Mutex::new(VecDeque::from(vec![
429 ScriptedResponse::Error(LlmError::Network("connection reset".into())),
430 ScriptedResponse::Text(fenced(GOOD_PLAN)),
431 ])),
432 });
433 let deps = make_deps(provider);
434
435 let outcome = drain(plan_and_run(
436 deps,
437 "flaky provider task",
438 CancellationToken::new(),
439 mpsc::channel(64).0,
440 ))
441 .await;
442
443 match outcome {
444 PlanOutcome::Completed { attempts, .. } => assert_eq!(attempts, 1),
445 other => panic!("expected completed after retry, got {other:?}"),
446 }
447}
448

1#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
2
3use std::collections::VecDeque;
4use std::sync::{Arc, Mutex};
5
6use tokio::sync::mpsc;
7use tokio_util::sync::CancellationToken;
8
9use tempfile::tempdir;
10
11use vak_agent::AutoApprove;
12use vak_flow::{
13 ExecutorDeps, PLANNER_SYSTEM, PlanOutcome, ToolCatalogEntry, build_planner_prompt,
14 extract_toml, plan_and_run, sanitize_basic_string_newlines,
15};
16use vak_llm::stream;
17use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, StopReason, Usage};
18use vak_llm::{EventStream, LlmError, Provider};
19use vak_permission::{Mode, PermissionEngine};
20use vak_tools::bash::BashTool;
21
22#[test]
23fn planner_prompt_includes_task_and_catalog() {
24 let catalog = vec![
25 ToolCatalogEntry {
26 name: "bash".into(),
27 description: "Execute a shell command".into(),
28 },
29 ToolCatalogEntry {
30 name: "task".into(),
31 description: "Delegate to a worker".into(),
32 },
33 ];
34 let prompt = build_planner_prompt("migrate the config module", &catalog);
35 assert!(prompt.contains("migrate the config module"));
36 assert!(prompt.contains("- bash: Execute a shell command"));
37 assert!(prompt.contains("- task: Delegate to a worker"));
38}
39
40#[test]
41fn extract_toml_handles_fenced_raw_and_garbage() {
42 let fenced = "Here is the plan:\n```toml\n[flow]\nname = \"x\"\n[[nodes]]\nid=\"a\"\ntype=\"bash\"\ncommand=\"echo hi\"\n```\ndone";
43 let extracted = extract_toml(fenced).unwrap();
44 assert!(extracted.contains("[flow]"));
45 assert!(!extracted.contains("```"));
46
47 let raw = "[flow]\nname = \"y\"\n\n[[nodes]]\nid = \"a\"\ntype = \"bash\"\ncommand = \"echo\"";
48 assert_eq!(extract_toml(raw).as_deref(), Some(raw));
49
50 assert!(extract_toml("I cannot plan this task, sorry.").is_none());
51}
52
53struct ScriptedPlanner {
54 /// Responses consumed in order across ALL requests (planner + children).
55 responses: Mutex<VecDeque<ScriptedResponse>>,
56}
57
58enum ScriptedResponse {
59 Text(String),
60 Error(LlmError),
61}
62
63#[async_trait::async_trait]
64impl Provider for ScriptedPlanner {
65 fn name(&self) -> &str {
66 "scripted"
67 }
68
69 async fn stream(
70 &self,
71 _request: ChatRequest,
72 _cancel: CancellationToken,
73 ) -> Result<EventStream, LlmError> {
74 let next = self.responses.lock().unwrap().pop_front();
75 let (mut sink, rx) = stream::channel(64);
76 match next {
77 Some(ScriptedResponse::Text(t)) => {
78 let msg = AssistantMessage {
79 content: vec![ContentBlock::text(t.clone())],
80 stop_reason: StopReason::EndTurn,
81 usage: Usage::default(),
82 model: "test-model".into(),
83 response_id: None,
84 };
85 sink.push(stream::StreamEvent::Start {
86 partial: msg.clone(),
87 });
88 sink.close_message(msg).await;
89 }
90 Some(ScriptedResponse::Error(e)) => sink.close_error(e).await,
91 None => sink.close_error(LlmError::Parse("exhausted".into())).await,
92 }
93 Ok(rx)
94 }
95}
96
97const GOOD_PLAN: &str = "\
98[flow]
99name = \"planned\"
100
101[[nodes]]
102id = \"probe\"
103type = \"bash\"
104command = \"echo planned-ok\"
105
106[[nodes]]
107id = \"report\"
108type = \"merge\"
109deps = [\"probe\"]
110";
111
112fn fenced(plan: &str) -> String {
113 format!("Sure, here is the plan:\n```toml\n{plan}\n```\n")
114}
115
116fn make_deps(provider: Arc<ScriptedPlanner>) -> Arc<ExecutorDeps> {
117 let dir = tempdir().unwrap();
118 let home = dir.path().join("home");
119 std::fs::create_dir_all(&home).unwrap();
120 std::mem::forget(dir);
121 Arc::new(ExecutorDeps {
122 prompt_layers: Vec::new(),
123 provider,
124 system_prompt: "sys".into(),
125 model: "test-model".into(),
126 tools: vec![Arc::new(BashTool)],
127 read_only_tools: vec![],
128 max_turns: 4,
129 outcome: None,
130 max_retries: 0,
131 retry_base_backoff_ms: 100,
132 request_timeout: Some(std::time::Duration::from_secs(600)),
133 circuit_breaker: None,
134 run_retry_attempts: 0,
135 run_retry_base_backoff_ms: 1000,
136 dispatch_ceiling: 1,
137 spend_gate: None,
138 permission: Some(Arc::new(PermissionEngine::default())),
139 mode: Mode::FullAccess,
140 approval_mode: vak_agent::ApprovalMode::Ask,
141 approver: Some(Arc::new(AutoApprove)),
142 sandbox: None,
143 cwd: std::env::temp_dir(),
144 sessions_home: home.clone(),
145 parent_session_id: "plan-parent".into(),
146 state_path: home.join("flow-runs/plan"),
147 agent_identity: None,
148 conversation_context: None,
149 work: None,
150 })
151}
152
153async fn drain(f: impl std::future::Future<Output = PlanOutcome>) -> PlanOutcome {
154 let (tx, mut rx) = mpsc::channel::<String>(256);
155 let d = tokio::spawn(async move { while rx.recv().await.is_some() {} });
156 drop(tx);
157 let out = f.await;
158 d.await.unwrap();
159 out
160}
161
162#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
163async fn valid_plan_executes_to_completion() {
164 let provider = Arc::new(ScriptedPlanner {
165 responses: Mutex::new(VecDeque::from(vec![ScriptedResponse::Text(fenced(
166 GOOD_PLAN,
167 ))])),
168 });
169 let deps = make_deps(provider);
170
171 let outcome = drain(plan_and_run(
172 deps,
173 "run the planned probe",
174 CancellationToken::new(),
175 mpsc::channel(64).0,
176 ))
177 .await;
178
179 match outcome {
180 PlanOutcome::Completed { outputs, attempts } => {
181 assert_eq!(attempts, 1);
182 assert!(
183 outputs
184 .get("probe")
185 .is_some_and(|o| o.contains("planned-ok"))
186 );
187 }
188 other => panic!("expected completed, got {other:?}"),
189 }
190}
191
192#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
193async fn garbage_plan_fails_closed_without_execution() {
194 let provider = Arc::new(ScriptedPlanner {
195 responses: Mutex::new(VecDeque::from(vec![ScriptedResponse::Text(
196 "I'm sorry, I cannot produce a TOML plan for that.".into(),
197 )])),
198 });
199 let deps = make_deps(provider);
200
201 let outcome = drain(plan_and_run(
202 deps,
203 "impossible task",
204 CancellationToken::new(),
205 mpsc::channel(64).0,
206 ))
207 .await;
208
209 match outcome {
210 PlanOutcome::PlanningFailed { reason } => {
211 assert!(reason.contains("no TOML plan"));
212 }
213 other => panic!("expected planning_failed, got {other:?}"),
214 }
215}
216
217#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
218async fn structurally_invalid_plan_fails_closed() {
219 // A cycle must be rejected by validation, never executed.
220 let cyclic = "\
221[flow]
222name = \"cyclic\"
223
224[[nodes]]
225id = \"a\"
226type = \"bash\"
227command = \"echo a\"
228deps = [\"b\"]
229
230[[nodes]]
231id = \"b\"
232type = \"bash\"
233command = \"echo b\"
234deps = [\"a\"]
235";
236 let provider = Arc::new(ScriptedPlanner {
237 responses: Mutex::new(VecDeque::from([ScriptedResponse::Text(fenced(cyclic))])),
238 });
239 let deps = make_deps(provider);
240
241 let outcome = drain(plan_and_run(
242 deps,
243 "circular task",
244 CancellationToken::new(),
245 mpsc::channel(64).0,
246 ))
247 .await;
248
249 match outcome {
250 PlanOutcome::PlanningFailed { reason } => {
251 assert!(reason.contains("cycle"), "got: {reason}");
252 }
253 other => panic!("expected planning_failed, got {other:?}"),
254 }
255}
256
257#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
258async fn failed_execution_triggers_exactly_one_replan() {
259 // Attempt 1: plan with a required node that fails at runtime.
260 // Attempt 2 (replan): clean plan that succeeds.
261 let failing_plan = "\
262[flow]
263name = \"attempt-one\"
264
265[[nodes]]
266id = \"boom\"
267type = \"bash\"
268command = \"exit 5\"
269required = true
270
271[[nodes]]
272id = \"report\"
273type = \"merge\"
274deps = [\"boom\"]
275";
276 let provider = Arc::new(ScriptedPlanner {
277 responses: Mutex::new(VecDeque::from(vec![
278 ScriptedResponse::Text(fenced(failing_plan)),
279 ScriptedResponse::Text(fenced(GOOD_PLAN)),
280 ])),
281 });
282 let deps = make_deps(provider);
283
284 let outcome = drain(plan_and_run(
285 deps,
286 "flaky task",
287 CancellationToken::new(),
288 mpsc::channel(256).0,
289 ))
290 .await;
291
292 match outcome {
293 PlanOutcome::Completed { attempts, .. } => {
294 assert_eq!(attempts, 2, "exactly one replan allowed");
295 }
296 other => panic!("expected completed after replan, got {other:?}"),
297 }
298}
299
300#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
301async fn replan_budget_is_bounded_at_one_retry() {
302 let bad_plan = "\
303[flow]
304name = \"always-fails\"
305
306[[nodes]]
307id = \"boom\"
308type = \"bash\"
309command = \"exit 7\"
310required = true
311";
312 let provider = Arc::new(ScriptedPlanner {
313 responses: Mutex::new(VecDeque::from(vec![
314 ScriptedResponse::Text(fenced(bad_plan)),
315 ScriptedResponse::Text(fenced(bad_plan)),
316 // A third plan would be accepted here if the budget were unbounded.
317 ScriptedResponse::Text(fenced(GOOD_PLAN)),
318 ])),
319 });
320 let deps = make_deps(provider);
321
322 let outcome = drain(plan_and_run(
323 deps,
324 "doomed task",
325 CancellationToken::new(),
326 mpsc::channel(256).0,
327 ))
328 .await;
329
330 match outcome {
331 PlanOutcome::Failed { node, .. } => {
332 assert_eq!(node, "boom");
333 }
334 other => panic!("expected failed after budget exhausted, got {other:?}"),
335 }
336}
337
338#[test]
339fn planner_system_prompt_is_compact() {
340 // Prompt discipline: the planner system prompt stays small.
341 let tokens_estimate = PLANNER_SYSTEM.len() / 4;
342 assert!(
343 tokens_estimate < 500,
344 "planner system prompt too large (~{tokens_estimate} tokens)"
345 );
346}
347
348#[test]
349fn sanitizer_escapes_newlines_in_basic_strings_only() {
350 let doc = "[flow]\nname = \"x\"\n\n[[nodes]]\nid = \"a\"\ntype = \"bash\"\ncommand = \"echo one\necho two\"\n";
351 let fixed = sanitize_basic_string_newlines(doc);
352 assert_eq!(
353 fixed,
354 "[flow]\nname = \"x\"\n\n[[nodes]]\nid = \"a\"\ntype = \"bash\"\ncommand = \"echo one\\necho two\"\n"
355 );
356 assert!(vak_flow::parse_flow(&fixed).is_ok());
357
358 // Valid documents pass through byte-for-byte.
359 assert_eq!(sanitize_basic_string_newlines(GOOD_PLAN), GOOD_PLAN);
360
361 // Multi-line basic strings are untouched (already valid TOML).
362 let multi = "prompt = \"\"\"a\nb\"\"\"\n";
363 assert_eq!(sanitize_basic_string_newlines(multi), multi);
364
365 // Literal strings and comments are untouched.
366 let lit = "command = 'echo hi'\n# comment with \" quote\nname = \"z\"\n";
367 assert_eq!(sanitize_basic_string_newlines(lit), lit);
368
369 // Escaped quotes/backslashes inside basic strings survive.
370 let escaped = "prompt = \"said \\\"hi\\\" then \\\\ broke\ninto two\"\n";
371 assert_eq!(
372 sanitize_basic_string_newlines(escaped),
373 "prompt = \"said \\\"hi\\\" then \\\\ broke\\ninto two\"\n"
374 );
375}
376
377#[test]
378fn sanitizer_escapes_nested_quotes_but_keeps_closers() {
379 // Nested shell quotes are escaped; real closers survive.
380 let doc = "[flow]\nname = \"q\"\n\n[[nodes]]\nid = \"a\"\ntype = \"bash\"\ncommand = \"test -z \"$(grep x f)\" && echo \"done\"\"\n";
381 let fixed = sanitize_basic_string_newlines(doc);
382 assert!(
383 vak_flow::parse_flow(&fixed).is_ok(),
384 "sanitized doc must parse: {fixed}"
385 );
386 assert!(fixed.contains("\\\"$(grep x f)\\\""));
387 assert!(fixed.contains("\\\"done\\\""));
388
389 // Closing quotes followed by comment / whitespace+comma stay closers.
390 let ok = "name = \"x\" # trailing\nnodes-note = [ \"a\" , \"b\" ]\n";
391 assert_eq!(sanitize_basic_string_newlines(ok), ok);
392}
393
394#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
395async fn multiline_basic_string_plan_is_sanitized_and_executes() {
396 // Regression: live planner runs emitted raw newlines inside
397 // `prompt = "..."`, which is invalid TOML. The structural sanitizer must
398 // repair it so the plan parses and executes instead of failing closed.
399 let plan = "[flow]\nname = \"multiline\"\n\n[[nodes]]\nid = \"probe\"\ntype = \"bash\"\ncommand = \"echo line-one\necho line-two\"\n\n[[nodes]]\nid = \"report\"\ntype = \"merge\"\ndeps = [\"probe\"]\n";
400 let provider = Arc::new(ScriptedPlanner {
401 responses: Mutex::new(VecDeque::from([ScriptedResponse::Text(fenced(plan))])),
402 });
403 let deps = make_deps(provider);
404
405 let outcome = drain(plan_and_run(
406 deps,
407 "multiline task",
408 CancellationToken::new(),
409 mpsc::channel(64).0,
410 ))
411 .await;
412
413 match outcome {
414 PlanOutcome::Completed { outputs, .. } => {
415 let out = outputs.get("probe").unwrap();
416 assert!(out.contains("line-one") && out.contains("line-two"));
417 }
418 other => panic!("expected completed after sanitize, got {other:?}"),
419 }
420}
421
422#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
423async fn transient_planner_failure_is_retried() {
424 // Invariant 7: the planner bypasses the loop's retry machinery, so it
425 // must retry transient provider failures itself instead of failing the
426 // whole run closed.
427 let provider = Arc::new(ScriptedPlanner {
428 responses: Mutex::new(VecDeque::from(vec![
429 ScriptedResponse::Error(LlmError::Network("connection reset".into())),
430 ScriptedResponse::Text(fenced(GOOD_PLAN)),
431 ])),
432 });
433 let deps = make_deps(provider);
434
435 let outcome = drain(plan_and_run(
436 deps,
437 "flaky provider task",
438 CancellationToken::new(),
439 mpsc::channel(64).0,
440 ))
441 .await;
442
443 match outcome {
444 PlanOutcome::Completed { attempts, .. } => assert_eq!(attempts, 1),
445 other => panic!("expected completed after retry, got {other:?}"),
446 }
447}
448

Indexing the workspace…