Vakyartha documentation

1#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
2
3use std::collections::VecDeque;
4use std::sync::{Arc, Mutex};
5use std::time::Instant;
6
7use tokio::sync::mpsc;
8use tokio_util::sync::CancellationToken;
9
10use tempfile::tempdir;
11
12use vak_agent::{Agent, AgentConfig, TaskDeps, TaskTool, TurnOutcome, WorkerRegistry};
13use vak_llm::stream;
14use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, StopReason, Usage};
15use vak_llm::{EventStream, LlmError, Provider};
16use vak_permission::PermissionEngine;
17use vak_session::types::{FrozenContract, SessionHeader};
18use vak_session::{SessionLog, SessionPath};
19use vak_tools::bash::BashTool;
20
21/// Routes scripted responses by the trailing user-message tag so parallel
22/// children are deterministic regardless of request interleaving.
23struct TaggedScripted {
24 routes: Mutex<HashMap<String, VecDeque<AssistantMessage>>>,
25}
26
27use std::collections::HashMap;
28
29impl TaggedScripted {
30 fn route_for(request: &ChatRequest) -> String {
31 request
32 .messages
33 .iter()
34 .find(|m| m.role == vak_llm::types::Role::User)
35 .map(|m| m.text_content())
36 .unwrap_or_default()
37 }
38}
39
40#[async_trait::async_trait]
41impl Provider for TaggedScripted {
42 fn name(&self) -> &str {
43 "scripted"
44 }
45
46 async fn stream(
47 &self,
48 request: ChatRequest,
49 _cancel: CancellationToken,
50 ) -> Result<EventStream, LlmError> {
51 let key = Self::route_for(&request);
52 let next = self
53 .routes
54 .lock()
55 .unwrap()
56 .get_mut(&key)
57 .and_then(|d| d.pop_front());
58 eprintln!(
59 "[route {key}] t={:?} -> {}",
60 Instant::now(),
61 match &next {
62 Some(m) => format!("{:?}", m.stop_reason),
63 None => "EXHAUSTED".into(),
64 }
65 );
66 let (mut sink, rx) = stream::channel(64);
67 match next {
68 Some(m) => {
69 sink.push(stream::StreamEvent::Start { partial: m.clone() });
70 sink.close_message(m).await;
71 }
72 None => {
73 sink.close_error(LlmError::Parse(format!("exhausted: {key}")))
74 .await
75 }
76 }
77 Ok(rx)
78 }
79}
80
81fn text(t: &str) -> AssistantMessage {
82 AssistantMessage {
83 content: vec![ContentBlock::text(t)],
84 stop_reason: StopReason::EndTurn,
85 usage: Usage::default(),
86 model: "test-model".into(),
87 response_id: None,
88 }
89}
90
91fn bash_script(command: String) -> AssistantMessage {
92 AssistantMessage {
93 content: vec![ContentBlock::ToolUse {
94 id: format!("s{}", uuid_like()),
95 name: "bash".into(),
96 input: serde_json::json!({ "command": command }),
97 }],
98 stop_reason: StopReason::ToolUse,
99 usage: Usage::default(),
100 model: "test-model".into(),
101 response_id: None,
102 }
103}
104
105/// What a worker's shell does. The test proves ordering with marker files, not
106/// with a stopwatch: elapsed time is a proxy that fails whenever the machine
107/// is busy, and the claim is about who ran alongside whom.
108///
109/// * `Rendezvous` — announce start, then wait (bounded) for the peer's start
110/// marker. It can only succeed if the two workers really ran at the same
111/// time; if the scheduler had serialized them the wait times out.
112/// * `AfterWave` — record whether both wave-1 workers had already finished when
113/// this one began.
114///
115/// Every script ends in `true`: recording an observation must never make the
116/// command itself fail, or the stop gate would add a turn and the test would
117/// be measuring that instead.
118enum Worker<'a> {
119 Rendezvous { me: &'a str, peer: &'a str },
120 AfterWave { me: &'a str },
121}
122
123fn worker_script(dir: &std::path::Path, worker: Worker<'_>) -> String {
124 let d = dir.display();
125 let body = match worker {
126 Worker::Rendezvous { me, peer } => format!(
127 "touch {d}/{me}.start; \
128 for i in $(seq 1 400); do [ -f {d}/{peer}.start ] && break; sleep 0.05; done; \
129 [ -f {d}/{peer}.start ] && touch {d}/{me}.saw_peer; \
130 sleep 0.3; touch {d}/{me}.end; true"
131 ),
132 Worker::AfterWave { me } => format!(
133 "touch {d}/{me}.start; \
134 [ -f {d}/a.end ] && [ -f {d}/b.end ] && touch {d}/{me}.saw_wave_done; true"
135 ),
136 };
137 format!("sh -c '{body}'")
138}
139
140fn task_call(id: &str, paths: &[&str], tag: &str) -> AssistantMessage {
141 AssistantMessage {
142 content: vec![ContentBlock::ToolUse {
143 id: id.into(),
144 name: "task".into(),
145 input: serde_json::json!({
146 "prompt": format!("do the thing for {tag}"),
147 "paths": paths,
148 }),
149 }],
150 stop_reason: StopReason::ToolUse,
151 usage: Usage::default(),
152 model: "test-model".into(),
153 response_id: None,
154 }
155}
156
157fn multi_task_msg(calls: Vec<AssistantMessage>) -> AssistantMessage {
158 let merged = calls
159 .into_iter()
160 .filter_map(|m| m.content.into_iter().next())
161 .collect();
162 AssistantMessage {
163 content: merged,
164 stop_reason: StopReason::ToolUse,
165 usage: Usage::default(),
166 model: "test-model".into(),
167 response_id: None,
168 }
169}
170
171fn uuid_like() -> String {
172 use std::sync::atomic::{AtomicU32, Ordering};
173 static C: AtomicU32 = AtomicU32::new(0);
174 format!("u{}", C.fetch_add(1, Ordering::Relaxed))
175}
176
177fn child_script(command: String) -> VecDeque<AssistantMessage> {
178 VecDeque::from(vec![bash_script(command), text("child done")])
179}
180
181#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
182async fn disjoint_writers_run_parallel_conflicting_writer_serializes() {
183 let dir = tempdir().unwrap();
184 let home = dir.path().join("home");
185 std::fs::create_dir_all(&home).unwrap();
186 let header = SessionHeader {
187 agent: None,
188 session_id: "fanout-parent".into(),
189 created_at: chrono::Utc::now(),
190 cwd: dir.path().to_path_buf(),
191 parent_session_id: None,
192 contract_id: None,
193 work_item_id: None,
194 conversation: None,
195 contract: FrozenContract {
196 app_version: "0".into(),
197 provider: "scripted".into(),
198 model: "test-model".into(),
199 route_ladder: Vec::new(),
200 route_objective: String::new(),
201 route_annotations: Vec::new(),
202 system_prompt: "sys".into(),
203 permission_mode: "full-access".into(),
204 capabilities: Vec::new(),
205 prompt_layers: Vec::new(),
206 },
207 };
208 let log = SessionLog::create(
209 SessionPath::new_session_file(&home, dir.path(), "fanout-parent"),
210 header,
211 )
212 .unwrap();
213
214 let mut routes: HashMap<String, VecDeque<AssistantMessage>> = HashMap::new();
215 routes.insert(
216 "fan out".into(),
217 VecDeque::from(vec![
218 multi_task_msg(vec![
219 task_call("a", &["src/a/**"], "A"),
220 task_call("b", &["src/b/**"], "B"),
221 task_call("c", &["src/a/**"], "C"),
222 ]),
223 text("all done"),
224 ]),
225 );
226 let work = dir.path().to_path_buf();
227 routes.insert(
228 "do the thing for A".into(),
229 child_script(worker_script(
230 &work,
231 Worker::Rendezvous { me: "a", peer: "b" },
232 )),
233 );
234 routes.insert(
235 "do the thing for B".into(),
236 child_script(worker_script(
237 &work,
238 Worker::Rendezvous { me: "b", peer: "a" },
239 )),
240 );
241 routes.insert(
242 "do the thing for C".into(),
243 child_script(worker_script(&work, Worker::AfterWave { me: "c" })),
244 );
245
246 let provider = Arc::new(TaggedScripted {
247 routes: Mutex::new(routes),
248 });
249
250 let mut cfg = AgentConfig::new("sys");
251 cfg.model = "test-model".into();
252 cfg.tools = vec![Arc::new(TaskTool::new(TaskDeps {
253 parent_agent_identity: None,
254 role_prompts: Default::default(),
255 provider: provider.clone(),
256 system_prompt: "child-sys".into(),
257 tail: Default::default(),
258 model: "test-model".into(),
259 tools: vec![Arc::new(BashTool)],
260 capabilities: Vec::new(),
261 hooks: None,
262 revocation_check: None,
263 presentation_rebuild: None,
264 mcp_tool_index: None,
265 input_normalizer: None,
266 read_only_tools: vec![],
267 max_turns: 4,
268 outcome_objective: None,
269 outcome: None,
270 max_retries: 0,
271 retry_base_backoff_ms: 0,
272 request_timeout: None,
273 circuit_breaker: None,
274 run_retry_attempts: 0,
275 run_retry_base_backoff_ms: 0,
276 dispatch_ceiling: 1,
277 spend_gate: None,
278 permission: Some(Arc::new(PermissionEngine::default())),
279 mode: vak_permission::Mode::FullAccess,
280 approval_mode: vak_agent::ApprovalMode::Ask,
281 approver: Some(Arc::new(vak_agent::AutoApprove)),
282 sandbox: None,
283 cwd: dir.path().to_path_buf(),
284 sessions_home: home.clone(),
285 parent_session_id: "fanout-parent".into(),
286 contract_id: None,
287 work_item_id: None,
288 work_item_ids: vec![],
289 events: None,
290 registry: Some(Arc::new(WorkerRegistry::new())),
291 }))];
292 cfg.permission = Some(Arc::new(
293 PermissionEngine::from_rule_strings(&["+task".to_string(), "+Bash(sh *)".to_string()])
294 .unwrap(),
295 ));
296 cfg.approver = Some(Arc::new(vak_agent::AutoApprove));
297 let mut agent = Agent::new(provider, log, cfg);
298 std::mem::forget(dir);
299
300 let (ev_tx, mut ev_rx) = mpsc::channel(4096);
301 let drainer = tokio::spawn(async move { while ev_rx.recv().await.is_some() {} });
302
303 let outcome = agent
304 .run(
305 "fan out",
306 &Default::default(),
307 CancellationToken::new(),
308 ev_tx,
309 )
310 .await;
311 drop(drainer);
312
313 assert!(
314 matches!(outcome, TurnOutcome::Completed { .. }),
315 "got {outcome:?}"
316 );
317
318 // Wave 1: A and B have disjoint paths, so they must have run at the same
319 // time — each saw the other's start marker while it was still running.
320 assert!(
321 work.join("a.saw_peer").exists() && work.join("b.saw_peer").exists(),
322 "disjoint writers A and B must run in parallel (each should have seen the other start)"
323 );
324 // Wave 2: C's paths conflict with A's, so it must wait for the wave to
325 // finish — both wave-1 workers were already done when C began.
326 assert!(
327 work.join("c.saw_wave_done").exists(),
328 "the conflicting writer C must start only after wave 1 (A and B) has finished"
329 );
330
331 let session = agent.session.lock().await;
332 // Raw ledger: the closed turn's results are trace lines in the
333 // projection now (docs/design/68-context-engine.md §10); source order
334 // is a property of what was recorded.
335 let msgs = session.message_chain();
336 let results: Vec<&str> = msgs
337 .iter()
338 .flat_map(|(_, m)| m.content.iter())
339 .filter_map(|b| match b {
340 ContentBlock::ToolResult { tool_use_id, .. } => Some(tool_use_id.as_str()),
341 _ => None,
342 })
343 .collect();
344 assert_eq!(
345 results,
346 vec!["a", "b", "c"],
347 "source order preserved across waves"
348 );
349}
350

1#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
2
3use std::collections::VecDeque;
4use std::sync::{Arc, Mutex};
5use std::time::Instant;
6
7use tokio::sync::mpsc;
8use tokio_util::sync::CancellationToken;
9
10use tempfile::tempdir;
11
12use vak_agent::{Agent, AgentConfig, TaskDeps, TaskTool, TurnOutcome, WorkerRegistry};
13use vak_llm::stream;
14use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, StopReason, Usage};
15use vak_llm::{EventStream, LlmError, Provider};
16use vak_permission::PermissionEngine;
17use vak_session::types::{FrozenContract, SessionHeader};
18use vak_session::{SessionLog, SessionPath};
19use vak_tools::bash::BashTool;
20
21/// Routes scripted responses by the trailing user-message tag so parallel
22/// children are deterministic regardless of request interleaving.
23struct TaggedScripted {
24 routes: Mutex<HashMap<String, VecDeque<AssistantMessage>>>,
25}
26
27use std::collections::HashMap;
28
29impl TaggedScripted {
30 fn route_for(request: &ChatRequest) -> String {
31 request
32 .messages
33 .iter()
34 .find(|m| m.role == vak_llm::types::Role::User)
35 .map(|m| m.text_content())
36 .unwrap_or_default()
37 }
38}
39
40#[async_trait::async_trait]
41impl Provider for TaggedScripted {
42 fn name(&self) -> &str {
43 "scripted"
44 }
45
46 async fn stream(
47 &self,
48 request: ChatRequest,
49 _cancel: CancellationToken,
50 ) -> Result<EventStream, LlmError> {
51 let key = Self::route_for(&request);
52 let next = self
53 .routes
54 .lock()
55 .unwrap()
56 .get_mut(&key)
57 .and_then(|d| d.pop_front());
58 eprintln!(
59 "[route {key}] t={:?} -> {}",
60 Instant::now(),
61 match &next {
62 Some(m) => format!("{:?}", m.stop_reason),
63 None => "EXHAUSTED".into(),
64 }
65 );
66 let (mut sink, rx) = stream::channel(64);
67 match next {
68 Some(m) => {
69 sink.push(stream::StreamEvent::Start { partial: m.clone() });
70 sink.close_message(m).await;
71 }
72 None => {
73 sink.close_error(LlmError::Parse(format!("exhausted: {key}")))
74 .await
75 }
76 }
77 Ok(rx)
78 }
79}
80
81fn text(t: &str) -> AssistantMessage {
82 AssistantMessage {
83 content: vec![ContentBlock::text(t)],
84 stop_reason: StopReason::EndTurn,
85 usage: Usage::default(),
86 model: "test-model".into(),
87 response_id: None,
88 }
89}
90
91fn bash_script(command: String) -> AssistantMessage {
92 AssistantMessage {
93 content: vec![ContentBlock::ToolUse {
94 id: format!("s{}", uuid_like()),
95 name: "bash".into(),
96 input: serde_json::json!({ "command": command }),
97 }],
98 stop_reason: StopReason::ToolUse,
99 usage: Usage::default(),
100 model: "test-model".into(),
101 response_id: None,
102 }
103}
104
105/// What a worker's shell does. The test proves ordering with marker files, not
106/// with a stopwatch: elapsed time is a proxy that fails whenever the machine
107/// is busy, and the claim is about who ran alongside whom.
108///
109/// * `Rendezvous` — announce start, then wait (bounded) for the peer's start
110/// marker. It can only succeed if the two workers really ran at the same
111/// time; if the scheduler had serialized them the wait times out.
112/// * `AfterWave` — record whether both wave-1 workers had already finished when
113/// this one began.
114///
115/// Every script ends in `true`: recording an observation must never make the
116/// command itself fail, or the stop gate would add a turn and the test would
117/// be measuring that instead.
118enum Worker<'a> {
119 Rendezvous { me: &'a str, peer: &'a str },
120 AfterWave { me: &'a str },
121}
122
123fn worker_script(dir: &std::path::Path, worker: Worker<'_>) -> String {
124 let d = dir.display();
125 let body = match worker {
126 Worker::Rendezvous { me, peer } => format!(
127 "touch {d}/{me}.start; \
128 for i in $(seq 1 400); do [ -f {d}/{peer}.start ] && break; sleep 0.05; done; \
129 [ -f {d}/{peer}.start ] && touch {d}/{me}.saw_peer; \
130 sleep 0.3; touch {d}/{me}.end; true"
131 ),
132 Worker::AfterWave { me } => format!(
133 "touch {d}/{me}.start; \
134 [ -f {d}/a.end ] && [ -f {d}/b.end ] && touch {d}/{me}.saw_wave_done; true"
135 ),
136 };
137 format!("sh -c '{body}'")
138}
139
140fn task_call(id: &str, paths: &[&str], tag: &str) -> AssistantMessage {
141 AssistantMessage {
142 content: vec![ContentBlock::ToolUse {
143 id: id.into(),
144 name: "task".into(),
145 input: serde_json::json!({
146 "prompt": format!("do the thing for {tag}"),
147 "paths": paths,
148 }),
149 }],
150 stop_reason: StopReason::ToolUse,
151 usage: Usage::default(),
152 model: "test-model".into(),
153 response_id: None,
154 }
155}
156
157fn multi_task_msg(calls: Vec<AssistantMessage>) -> AssistantMessage {
158 let merged = calls
159 .into_iter()
160 .filter_map(|m| m.content.into_iter().next())
161 .collect();
162 AssistantMessage {
163 content: merged,
164 stop_reason: StopReason::ToolUse,
165 usage: Usage::default(),
166 model: "test-model".into(),
167 response_id: None,
168 }
169}
170
171fn uuid_like() -> String {
172 use std::sync::atomic::{AtomicU32, Ordering};
173 static C: AtomicU32 = AtomicU32::new(0);
174 format!("u{}", C.fetch_add(1, Ordering::Relaxed))
175}
176
177fn child_script(command: String) -> VecDeque<AssistantMessage> {
178 VecDeque::from(vec![bash_script(command), text("child done")])
179}
180
181#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
182async fn disjoint_writers_run_parallel_conflicting_writer_serializes() {
183 let dir = tempdir().unwrap();
184 let home = dir.path().join("home");
185 std::fs::create_dir_all(&home).unwrap();
186 let header = SessionHeader {
187 agent: None,
188 session_id: "fanout-parent".into(),
189 created_at: chrono::Utc::now(),
190 cwd: dir.path().to_path_buf(),
191 parent_session_id: None,
192 contract_id: None,
193 work_item_id: None,
194 conversation: None,
195 contract: FrozenContract {
196 app_version: "0".into(),
197 provider: "scripted".into(),
198 model: "test-model".into(),
199 route_ladder: Vec::new(),
200 route_objective: String::new(),
201 route_annotations: Vec::new(),
202 system_prompt: "sys".into(),
203 permission_mode: "full-access".into(),
204 capabilities: Vec::new(),
205 prompt_layers: Vec::new(),
206 },
207 };
208 let log = SessionLog::create(
209 SessionPath::new_session_file(&home, dir.path(), "fanout-parent"),
210 header,
211 )
212 .unwrap();
213
214 let mut routes: HashMap<String, VecDeque<AssistantMessage>> = HashMap::new();
215 routes.insert(
216 "fan out".into(),
217 VecDeque::from(vec![
218 multi_task_msg(vec![
219 task_call("a", &["src/a/**"], "A"),
220 task_call("b", &["src/b/**"], "B"),
221 task_call("c", &["src/a/**"], "C"),
222 ]),
223 text("all done"),
224 ]),
225 );
226 let work = dir.path().to_path_buf();
227 routes.insert(
228 "do the thing for A".into(),
229 child_script(worker_script(
230 &work,
231 Worker::Rendezvous { me: "a", peer: "b" },
232 )),
233 );
234 routes.insert(
235 "do the thing for B".into(),
236 child_script(worker_script(
237 &work,
238 Worker::Rendezvous { me: "b", peer: "a" },
239 )),
240 );
241 routes.insert(
242 "do the thing for C".into(),
243 child_script(worker_script(&work, Worker::AfterWave { me: "c" })),
244 );
245
246 let provider = Arc::new(TaggedScripted {
247 routes: Mutex::new(routes),
248 });
249
250 let mut cfg = AgentConfig::new("sys");
251 cfg.model = "test-model".into();
252 cfg.tools = vec![Arc::new(TaskTool::new(TaskDeps {
253 parent_agent_identity: None,
254 role_prompts: Default::default(),
255 provider: provider.clone(),
256 system_prompt: "child-sys".into(),
257 tail: Default::default(),
258 model: "test-model".into(),
259 tools: vec![Arc::new(BashTool)],
260 capabilities: Vec::new(),
261 hooks: None,
262 revocation_check: None,
263 presentation_rebuild: None,
264 mcp_tool_index: None,
265 input_normalizer: None,
266 read_only_tools: vec![],
267 max_turns: 4,
268 outcome_objective: None,
269 outcome: None,
270 max_retries: 0,
271 retry_base_backoff_ms: 0,
272 request_timeout: None,
273 circuit_breaker: None,
274 run_retry_attempts: 0,
275 run_retry_base_backoff_ms: 0,
276 dispatch_ceiling: 1,
277 spend_gate: None,
278 permission: Some(Arc::new(PermissionEngine::default())),
279 mode: vak_permission::Mode::FullAccess,
280 approval_mode: vak_agent::ApprovalMode::Ask,
281 approver: Some(Arc::new(vak_agent::AutoApprove)),
282 sandbox: None,
283 cwd: dir.path().to_path_buf(),
284 sessions_home: home.clone(),
285 parent_session_id: "fanout-parent".into(),
286 contract_id: None,
287 work_item_id: None,
288 work_item_ids: vec![],
289 events: None,
290 registry: Some(Arc::new(WorkerRegistry::new())),
291 }))];
292 cfg.permission = Some(Arc::new(
293 PermissionEngine::from_rule_strings(&["+task".to_string(), "+Bash(sh *)".to_string()])
294 .unwrap(),
295 ));
296 cfg.approver = Some(Arc::new(vak_agent::AutoApprove));
297 let mut agent = Agent::new(provider, log, cfg);
298 std::mem::forget(dir);
299
300 let (ev_tx, mut ev_rx) = mpsc::channel(4096);
301 let drainer = tokio::spawn(async move { while ev_rx.recv().await.is_some() {} });
302
303 let outcome = agent
304 .run(
305 "fan out",
306 &Default::default(),
307 CancellationToken::new(),
308 ev_tx,
309 )
310 .await;
311 drop(drainer);
312
313 assert!(
314 matches!(outcome, TurnOutcome::Completed { .. }),
315 "got {outcome:?}"
316 );
317
318 // Wave 1: A and B have disjoint paths, so they must have run at the same
319 // time — each saw the other's start marker while it was still running.
320 assert!(
321 work.join("a.saw_peer").exists() && work.join("b.saw_peer").exists(),
322 "disjoint writers A and B must run in parallel (each should have seen the other start)"
323 );
324 // Wave 2: C's paths conflict with A's, so it must wait for the wave to
325 // finish — both wave-1 workers were already done when C began.
326 assert!(
327 work.join("c.saw_wave_done").exists(),
328 "the conflicting writer C must start only after wave 1 (A and B) has finished"
329 );
330
331 let session = agent.session.lock().await;
332 // Raw ledger: the closed turn's results are trace lines in the
333 // projection now (docs/design/68-context-engine.md §10); source order
334 // is a property of what was recorded.
335 let msgs = session.message_chain();
336 let results: Vec<&str> = msgs
337 .iter()
338 .flat_map(|(_, m)| m.content.iter())
339 .filter_map(|b| match b {
340 ContentBlock::ToolResult { tool_use_id, .. } => Some(tool_use_id.as_str()),
341 _ => None,
342 })
343 .collect();
344 assert_eq!(
345 results,
346 vec!["a", "b", "c"],
347 "source order preserved across waves"
348 );
349}
350

Fanout

Indexing the workspace…

Fanout