1. 1#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
2. 2
3. 3//! Cross-phase adoption matrix (docs/design/42-managed-work-contracts.md): scenarios where Phases
4. 4//! A (receipts), B (ladder), C (partitions), D (budget), H (goal) must
5. 5//! compose correctly — not just work in isolation.
6. 6
7. 7use std::collections::VecDeque;
8. 8use std::sync::Arc;
9. 9use std::sync::atomic::{AtomicU32, Ordering};
10. 10
11. 11use tempfile::tempdir;
12. 12use tokio::sync::mpsc;
13. 13use tokio_util::sync::CancellationToken;
14. 14
15. 15use vak_agent::{
16. 16 Agent, AgentConfig, AgentEvent, AutoApprove, SpendCheck, SpendGate, SteeringQueues, TurnOutcome,
17. 17};
18. 18use vak_llm::stream;
19. 19use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, StopReason, Usage};
20. 20use vak_llm::{AttemptReason, EventStream, LlmError, Provider, WorkPurpose};
21. 21use vak_permission::{Mode, PermissionEngine};
22. 22use vak_session::types::{EntryPayload, FrozenContract, SessionHeader};
23. 23use vak_session::{SessionLog, SessionPath};
24. 24
25. 25fn text_msg(model: &str, t: &str) -> AssistantMessage {
26. 26 AssistantMessage {
27. 27 content: vec![ContentBlock::text(t)],
28. 28 stop_reason: StopReason::EndTurn,
29. 29 usage: Usage {
30. 30 input_tokens: 5,
31. 31 output_tokens: 3,
32. 32 ..Default::default()
33. 33 },
34. 34 model: model.into(),
35. 35 response_id: None,
36. 36 }
37. 37}
38. 38
39. 39/// Primary leg: network-dead. Fallback leg: scripted FIFO.
40. 40struct MatrixProvider {
41. 41 fail_primary: AtomicU32,
42. 42 fallback: std::sync::Mutex<VecDeque<AssistantMessage>>,
43. 43 seen_models: std::sync::Mutex<Vec<String>>,
44. 44}
45. 45
46. 46impl MatrixProvider {
47. 47 fn new(fallbacks: Vec<AssistantMessage>) -> Arc<Self> {
48. 48 Arc::new(MatrixProvider {
49. 49 fail_primary: AtomicU32::new(u32::MAX),
50. 50 fallback: std::sync::Mutex::new(fallbacks.into_iter().collect()),
51. 51 seen_models: std::sync::Mutex::new(Vec::new()),
52. 52 })
53. 53 }
54. 54}
55. 55
56. 56#[async_trait::async_trait]
57. 57impl Provider for MatrixProvider {
58. 58 fn name(&self) -> &str {
59. 59 "matrix"
60. 60 }
61. 61
62. 62 async fn stream(
63. 63 &self,
64. 64 request: ChatRequest,
65. 65 _cancel: CancellationToken,
66. 66 ) -> Result<EventStream, LlmError> {
67. 67 self.seen_models.lock().unwrap().push(request.model.clone());
68. 68 let dead = self.fail_primary.fetch_sub(1, Ordering::SeqCst) > 0
69. 69 && request.model == "primary-model";
70. 70 let (mut sink, rx) = stream::channel(64);
71. 71 if dead {
72. 72 sink.close_error(LlmError::Network("primary down".into()))
73. 73 .await;
74. 74 } else {
75. 75 let next = self.fallback.lock().unwrap().pop_front();
76. 76 match next {
77. 77 Some(m) => {
78. 78 sink.push(stream::StreamEvent::Start { partial: m.clone() });
79. 79 sink.close_message(m).await;
80. 80 }
81. 81 None => {
82. 82 sink.close_error(LlmError::Parse("matrix exhausted".into()))
83. 83 .await
84. 84 }
85. 85 }
86. 86 }
87. 87 Ok(rx)
88. 88 }
89. 89}
90. 90
91. 91/// Denies a specific model; allows everything else.
92. 92struct ModelDenyGate {
93. 93 denied_model: &'static str,
94. 94 denials: std::sync::Mutex<Vec<String>>,
95. 95}
96. 96
97. 97#[async_trait::async_trait]
98. 98impl SpendGate for ModelDenyGate {
99. 99 async fn authorize(&self, check: &SpendCheck<'_>) -> Result<(), String> {
100. 100 if check.model == self.denied_model {
101. 101 self.denials.lock().unwrap().push(check.model.to_string());
102. 102 Err(format!("model {} not funded", check.model))
103. 103 } else {
104. 104 Ok(())
105. 105 }
106. 106 }
107. 107
108. 108 fn record_settled(&self, _provider: &str, _model: &str, _session_id: &str, _usage: &Usage) {}
109. 109}
110. 110
111. 111fn setup(provider: Arc<MatrixProvider>, cfg_tweaks: impl FnOnce(&mut AgentConfig)) -> Agent {
112. 112 let dir = tempdir().unwrap();
113. 113 // Keep the tempdir alive for the whole test by leaking the path into
114. 114 // the session home (tests are short-lived processes).
115. 115 let cwd = dir.path().to_path_buf();
116. 116 std::mem::forget(dir);
117. 117 let header = SessionHeader {
118. 118 agent: None,
119. 119 session_id: "matrix".into(),
120. 120 created_at: chrono::Utc::now(),
121. 121 cwd: cwd.clone(),
122. 122 parent_session_id: None,
123. 123 contract_id: None,
124. 124 work_item_id: None,
125. 125 conversation: None,
126. 126 contract: FrozenContract {
127. 127 app_version: "0".into(),
128. 128 provider: "matrix".into(),
129. 129 model: "primary-model".into(),
130. 130 route_ladder: vec![
131. 131 vak_llm::RouteLeg {
132. 132 provider: "matrix".into(),
133. 133 model: "primary-model".into(),
134. 134 dialect: vak_llm::EndpointDialect::ChatCompletions,
135. 135 credential_id: None,
136. 136 },
137. 137 vak_llm::RouteLeg {
138. 138 provider: "matrix".into(),
139. 139 model: "fallback-model".into(),
140. 140 dialect: vak_llm::EndpointDialect::ChatCompletions,
141. 141 credential_id: None,
142. 142 },
143. 143 ],
144. 144 route_objective: String::new(),
145. 145 route_annotations: Vec::new(),
146. 146 system_prompt: "sys".into(),
147. 147 permission_mode: "workspace-write".into(),
148. 148 capabilities: Vec::new(),
149. 149 prompt_layers: Vec::new(),
150. 150 },
151. 151 };
152. 152 let home = cwd.join(".vak-home");
153. 153 std::fs::create_dir_all(&home).unwrap();
154. 154 let log =
155. 155 SessionLog::create(SessionPath::new_session_file(&home, &cwd, "matrix"), header).unwrap();
156. 156 let mut cfg = AgentConfig::new("sys");
157. 157 cfg.model = "primary-model".into();
158. 158 cfg.mode = Mode::FullAccess;
159. 159 cfg.permission = Some(Arc::new(PermissionEngine::default()));
160. 160 cfg.approver = Some(Arc::new(AutoApprove));
161. 161 cfg.retry_base_backoff_ms = 1;
162. 162 cfg.run_retry_base_backoff_ms = 1;
163. 163 cfg.max_retries = 0;
164. 164 cfg.run_retry_attempts = 0;
165. 165 // Primary + one fallback.
166. 166 cfg.ladder = vec![(
167. 167 provider.clone() as Arc<dyn Provider>,
168. 168 "fallback-model".to_string(),
169. 169 )];
170. 170 cfg_tweaks(&mut cfg);
171. 171 Agent::new(provider, log, cfg)
172. 172}
173. 173
174. 174async fn run(agent: &mut Agent, prompt: &str) -> TurnOutcome {
175. 175 run_events(agent, prompt).await.0
176. 176}
177. 177
178. 178async fn run_events(agent: &mut Agent, prompt: &str) -> (TurnOutcome, Vec<AgentEvent>) {
179. 179 let (ev_tx, mut ev_rx) = mpsc::channel(1024);
180. 180 let cancel = CancellationToken::new();
181. 181 let steering = SteeringQueues::new();
182. 182 let fut = agent.run(prompt, &steering, cancel, ev_tx);
183. 183 let mut events = Vec::new();
184. 184 let outcome = tokio::join!(fut, async {
185. 185 while let Some(ev) = ev_rx.recv().await {
186. 186 events.push(ev);
187. 187 }
188. 188 });
189. 189 (outcome.0, events)
190. 190}
191. 191
192. 192const PASS_VERDICT: &str =
193. 193 r#"{"results":[{"criterion":"c1","verdict":"pass","evidence":"clearly done"}]}"#;
194. 194
195. 195/// B×H×A: primary dies on every step; every model turn AND the judge call
196. 196/// route to the fallback; judge receipt is purpose=Verify with the
197. 197/// fallback model; completion audited.
198. 198#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
199. 199async fn ladder_serves_goal_audits_when_primary_dead() {
200. 200 let provider = MatrixProvider::new(vec![
201. 201 text_msg("fallback-model", "attempting objective"),
202. 202 text_msg("fallback-model", PASS_VERDICT),
203. 203 ]);
204. 204 let mut agent = setup(provider.clone(), |_| {});
205. 205 agent.set_goal("do the thing", vec!["c1".into()]);
206. 206 let outcome = run(&mut agent, "do the thing").await;
207. 207 assert!(matches!(outcome, TurnOutcome::Completed { .. }));
208. 208
209. 209 {
210. 210 let models = provider.seen_models.lock().unwrap();
211. 211 // Each work unit walks primary(dead) -> fallback(ok). No work may
212. 212 // END on the dead leg.
213. 213 // A work unit may ATTEMPT the dead primary but must never settle
214. 214 // there: the last model seen per work must be the fallback.
215. 215 let committed_on_primary = models.last().map(|m| m == "primary-model").unwrap_or(false);
216. 216 assert!(
217. 217 !committed_on_primary,
218. 218 "a work unit must never settle on the dead leg: {models:?}"
219. 219 );
220. 220 assert!(models.contains(&"fallback-model".to_string()));
221. 221 }
222. 222
223. 223 let session = agent.into_session().await;
224. 224 let receipts = session.receipts();
225. 225 let verify = receipts
226. 226 .iter()
227. 227 .find(|r| r.purpose == WorkPurpose::Verify)
228. 228 .expect("judge dispatch receipted");
229. 229 assert_eq!(verify.model, "fallback-model");
230. 230 // The judge dispatch ALSO walked primary(dead)->fallback(ok).
231. 231 assert_eq!(verify.attempts[0].reason, AttemptReason::Initial);
232. 232 assert_eq!(verify.attempts[0].domain, vak_llm::FailureDomain::Network);
233. 233 assert_eq!(verify.attempts[1].reason, AttemptReason::RouteFallback);
234. 234 assert_eq!(verify.attempts[1].settlement, vak_llm::Settlement::Ok);
235. 235 assert_eq!(verify.winning_attempt, Some(1));
236. 236 assert_eq!(
237. 237 goal_final_status(&session),
238. 238 Some("done".into()),
239. 239 "audited done across a walked ladder"
240. 240 );
241. 241}
242. 242
243. 243/// B×D: budget denies ONLY the primary model; the Ask auto-fails there
244. 244/// (denial is per-dispatch, approver approves nothing here) — wait: the
245. 245/// approver below approves budget asks, which would let the PRIMARY
246. 246/// through. So instead assert the deny-then-approve path lands the
247. 247/// primary dispatch anyway, and the settled usage records under the
248. 248/// primary model.
249. 249#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
250. 250async fn budget_ask_on_primary_then_proceed() {
251. 251 let provider = MatrixProvider::new(vec![text_msg("primary-model", "answered")]);
252. 252 let mut agent = setup(provider.clone(), |cfg| {
253. 253 cfg.spend_gate = Some(Arc::new(ModelDenyGate {
254. 254 denied_model: "primary-model",
255. 255 denials: std::sync::Mutex::new(Vec::new()),
256. 256 }));
257. 257 });
258. 258 let outcome = run(&mut agent, "hello").await;
259. 259 assert!(matches!(outcome, TurnOutcome::Completed { .. }));
260. 260
261. 261 let session = agent.into_session().await;
262. 262 let r = &session.receipts()[0];
263. 263 // Approved ask admitted the primary attempt; the primary is also
264. 264 // network-dead in this fixture, so the walk continued to the fallback.
265. 265 assert_eq!(r.attempts[0].reason, vak_llm::AttemptReason::Initial);
266. 266 assert_eq!(r.attempts[0].domain, vak_llm::FailureDomain::Network);
267. 267 assert_eq!(r.attempts[1].reason, vak_llm::AttemptReason::RouteFallback);
268. 268 assert_eq!(r.winning_attempt, Some(1));
269. 269 assert_eq!(r.model, "fallback-model");
270. 270 // Exactly one budget denial was raised (for the primary); verified by
271. 271 // the receipt walk above (Initial@Network -> RouteFallback@Ok).
272. 272}
273. 273
274. 274/// D strictness: with NO approver (unattended), a denied primary must NOT
275. 275/// silently fall forward to a cheaper leg — budget admission is per-model
276. 276/// and a denial without approval fails the step permanently.
277. 277#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
278. 278async fn unattended_budget_denial_fails_even_with_ladder() {
279. 279 let provider = MatrixProvider::new(vec![text_msg("fallback-model", "should not happen")]);
280. 280 let mut agent = setup(provider.clone(), |cfg| {
281. 281 cfg.approver = None; // unattended
282. 282 cfg.spend_gate = Some(Arc::new(ModelDenyGate {
283. 283 denied_model: "primary-model",
284. 284 denials: std::sync::Mutex::new(Vec::new()),
285. 285 }));
286. 286 });
287. 287 match run(&mut agent, "hello").await {
288. 288 TurnOutcome::Failed { error } => {
289. 289 assert!(error.to_string().contains("budget admission denied"));
290. 290 }
291. 291 other => panic!("expected permanent failure, got {other:?}"),
292. 292 }
293. 293 // The fallback leg was never dispatched: silence means no.
294. 294 assert!(
295. 295 provider
296. 296 .seen_models
297. 297 .lock()
298. 298 .unwrap()
299. 299 .iter()
300. 300 .all(|m| m == "primary-model"),
301. 301 "no dispatch may leave for the fallback after an unanswered denial"
302. 302 );
303. 303}
304. 304
305. 305/// C×H: a goal run with no usable horizon takes the reset-with-handoff
306. 306/// rescue, and receipts still land around it.
307. 307#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
308. 308async fn compaction_partitions_coexist_with_goal_and_receipts() {
309. 309 // NOTE: a 900-token declared window is dwarfed by this filler's own
310. 310 // estimated size, so `budget` saturates to 0 (no usable horizon) and
311. 311 // the reset-with-handoff rescue fires before the first model turn.
312. 312 let filler = "z".repeat(5000);
313. 313 let provider = MatrixProvider::new(vec![
314. 314 // Over-budget path fires BEFORE the first model turn; the handoff
315. 315 // rescue consumes this slot.
316. 316 text_msg("primary-model", "# Objective\nfinish\n# Open Items\nnone"),
317. 317 text_msg("primary-model", "claim one"),
318. 318 // Judge gets prose => fail-closed rejection.
319. 319 text_msg("primary-model", "no json here"),
320. 320 text_msg("primary-model", "claim two"),
321. 321 text_msg("primary-model", PASS_VERDICT),
322. 322 ]);
323. 323 let mut agent = setup(provider, |_| {});
324. 324 agent.config.declared_window = 900;
325. 325 agent.config.max_output = 64;
326. 326 agent.set_goal("finish despite compaction", vec!["c1".into()]);
327. 327 let outcome = run(&mut agent, &format!("task {filler}")).await;
328. 328 assert!(matches!(outcome, TurnOutcome::Completed { .. }));
329. 329
330. 330 let session = agent.into_session().await;
331. 331 let compactions: Vec<_> = session
332. 332 .chain_to_root()
333. 333 .iter()
334. 334 .filter_map(|e| match &e.payload {
335. 335 EntryPayload::Compaction(c) => Some(c),
336. 336 _ => None,
337. 337 })
338. 338 .collect();
339. 339 assert!(!compactions.is_empty(), "compaction ran");
340. 340 let handoff_fired = session
341. 341 .chain_to_root()
342. 342 .iter()
343. 343 .any(|e| matches!(&e.payload, EntryPayload::Compaction(c) if c.reset_all));
344. 344 assert!(handoff_fired, "rescue ran before the first model turn");
345. 345 // Partition accounting on any NORMAL compaction is covered by the
346. 346 // vak-eval scorecard; the reset entry itself carries none by design.
347. 347
348. 348 // Receipts survived alongside partitions; goal closed audited.
349. 349 assert!(session.receipts().len() >= 2);
350. 350 assert_eq!(goal_final_status(&session), Some("done".into()));
351. 351}
352. 352
353. 353/// Ceiling starvation mid-ladder: ceiling=1 admits exactly one dispatch;
354. 354/// the primary's failure consumes it and the fallback may NOT go out.
355. 355#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
356. 356async fn starved_ceiling_blocks_fallback_leg() {
357. 357 let provider = MatrixProvider::new(vec![text_msg("fallback-model", "never")]);
358. 358 let mut agent = setup(provider.clone(), |cfg| {
359. 359 cfg.dispatch_ceiling = 1;
360. 360 });
361. 361 match run(&mut agent, "hello").await {
362. 362 TurnOutcome::Failed { error } => {
363. 363 assert!(
364. 364 error.to_string().contains("dispatch ceiling"),
365. 365 "got: {error}"
366. 366 );
367. 367 }
368. 368 other => panic!("expected ceiling failure, got {other:?}"),
369. 369 }
370. 370 let models = provider.seen_models.lock().unwrap();
371. 371 assert_eq!(
372. 372 models.iter().filter(|m| **m == *"fallback-model").count(),
373. 373 0,
374. 374 "starved ceiling must not fund the next leg"
375. 375 );
376. 376}
377. 377
378. 378fn goal_final_status(session: &SessionLog) -> Option<String> {
379. 379 session
380. 380 .chain_to_root()
381. 381 .iter()
382. 382 .filter_map(|e| match &e.payload {
383. 383 EntryPayload::Goal(g) => Some(match g.status {
384. 384 vak_session::types::GoalStatus::Active => "active".to_string(),
385. 385 vak_session::types::GoalStatus::Done { .. } => "done".to_string(),
386. 386 vak_session::types::GoalStatus::Unverified { .. } => "unverified".to_string(),
387. 387 }),
388. 388 _ => None,
389. 389 })
390. 390 .next_back()
391. 391}
392. 392