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, TurnOutcome};
13use vak_llm::stream;
14use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, StopReason, Usage};
15use vak_llm::{EventStream, LlmError, Provider};
16use vak_session::SessionLog;
17use vak_session::types::{FrozenContract, SessionHeader};
18
19enum Step {
20 Err(LlmError),
21 Text(String),
22}
23
24struct Scripted {
25 steps: Mutex<VecDeque<Step>>,
26}
27
28#[async_trait::async_trait]
29impl Provider for Scripted {
30 fn name(&self) -> &str {
31 "scripted"
32 }
33
34 async fn stream(
35 &self,
36 _request: ChatRequest,
37 _cancel: CancellationToken,
38 ) -> Result<EventStream, LlmError> {
39 let next = self.steps.lock().unwrap().pop_front();
40 let (mut sink, rx) = stream::channel(64);
41 match next {
42 Some(Step::Text(t)) => {
43 let msg = AssistantMessage {
44 content: vec![ContentBlock::text(t)],
45 stop_reason: StopReason::EndTurn,
46 usage: Usage::default(),
47 model: "test-model".into(),
48 response_id: None,
49 };
50 sink.push(stream::StreamEvent::Start {
51 partial: msg.clone(),
52 });
53 sink.close_message(msg).await;
54 }
55 Some(Step::Err(e)) => sink.close_error(e).await,
56 None => sink.close_error(LlmError::Parse("exhausted".into())).await,
57 }
58 Ok(rx)
59 }
60}
61
62fn build(
63 steps: Vec<Step>,
64 max_retries: u32,
65 base_ms: u64,
66 timeout: Option<std::time::Duration>,
67) -> Agent {
68 let dir = tempdir().unwrap();
69 let header = SessionHeader {
70 agent: None,
71 session_id: "rel".into(),
72 created_at: chrono::Utc::now(),
73 cwd: dir.path().to_path_buf(),
74 parent_session_id: None,
75 contract_id: None,
76 work_item_id: None,
77 conversation: None,
78 contract: FrozenContract {
79 app_version: "0".into(),
80 provider: "scripted".into(),
81 model: "test-model".into(),
82 route_ladder: Vec::new(),
83 route_objective: String::new(),
84 route_annotations: Vec::new(),
85 system_prompt: "sys".into(),
86 permission_mode: "full-access".into(),
87 capabilities: Vec::new(),
88 prompt_layers: Vec::new(),
89 },
90 };
91 let log = SessionLog::create(dir.path().join("s.jsonl"), header).unwrap();
92 let mut cfg = AgentConfig::new("sys");
93 cfg.max_retries = max_retries;
94 cfg.retry_base_backoff_ms = base_ms;
95 cfg.request_timeout = timeout;
96 // Step-level machinery is the system under test here; run-level
97 // endurance has its own tests (tests/run_endurance.rs).
98 cfg.run_retry_attempts = 0;
99 std::mem::forget(dir);
100 Agent::new(
101 Arc::new(Scripted {
102 steps: Mutex::new(steps.into_iter().collect()),
103 }),
104 log,
105 cfg,
106 )
107}
108
109#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
110async fn transient_errors_are_retried_until_success() {
111 let mut agent = build(
112 vec![
113 Step::Err(LlmError::RateLimit {
114 message: "slow down".into(),
115 retry_after_secs: None,
116 }),
117 Step::Err(LlmError::Overloaded("529".into())),
118 Step::Err(LlmError::Network("flap".into())),
119 Step::Text("finally".into()),
120 ],
121 3,
122 1, // near-zero backoff for the test
123 None,
124 );
125
126 let outcome = agent
127 .run(
128 "go",
129 &Default::default(),
130 CancellationToken::new(),
131 mpsc::channel(64).0,
132 )
133 .await;
134
135 match outcome {
136 TurnOutcome::Completed { response } => assert_eq!(response.text_content(), "finally"),
137 other => panic!("expected completed after retries, got {other:?}"),
138 }
139}
140
141#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
142async fn non_retryable_errors_fail_immediately() {
143 let mut agent = build(
144 vec![
145 Step::Err(LlmError::Auth("bad key".into())),
146 Step::Text("never".into()),
147 ],
148 3,
149 1,
150 None,
151 );
152
153 let start = Instant::now();
154 let outcome = agent
155 .run(
156 "go",
157 &Default::default(),
158 CancellationToken::new(),
159 mpsc::channel(64).0,
160 )
161 .await;
162 assert!(
163 start.elapsed() < std::time::Duration::from_millis(500),
164 "auth errors must not be retried"
165 );
166 match outcome {
167 TurnOutcome::Failed { error } => assert!(error.to_string().contains("bad key")),
168 other => panic!("expected failed, got {other:?}"),
169 }
170}
171
172#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
173async fn retry_budget_exhaustion_fails_with_last_error() {
174 let err_fn = || {
175 Step::Err(LlmError::RateLimit {
176 message: "still limited".into(),
177 retry_after_secs: None,
178 })
179 };
180 let mut agent = build(
181 vec![err_fn(), err_fn(), err_fn(), err_fn(), err_fn()],
182 2, // budget smaller than failure count
183 1,
184 None,
185 );
186
187 let outcome = agent
188 .run(
189 "go",
190 &Default::default(),
191 CancellationToken::new(),
192 mpsc::channel(64).0,
193 )
194 .await;
195 match outcome {
196 TurnOutcome::Failed { error } => assert!(error.to_string().contains("still limited")),
197 other => panic!("expected failed after budget, got {other:?}"),
198 }
199}
200
201#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
202async fn watchdog_deadline_converts_hung_step_into_retryable_failure() {
203 // A provider whose stream never completes: the deadline must fire.
204 let cancelled = Arc::new(std::sync::atomic::AtomicBool::new(false));
205 struct Hung(Arc<std::sync::atomic::AtomicBool>);
206 #[async_trait::async_trait]
207 impl Provider for Hung {
208 fn name(&self) -> &str {
209 "hung"
210 }
211 async fn stream(
212 &self,
213 _r: ChatRequest,
214 c: CancellationToken,
215 ) -> Result<EventStream, LlmError> {
216 let cancelled = self.0.clone();
217 let (mut sink, rx) = stream::channel(8);
218 sink.push(stream::StreamEvent::Start {
219 partial: AssistantMessage::empty("m"),
220 });
221 // Hold the sink open forever — simulates a stalled stream.
222 tokio::spawn(async move {
223 let _keep_alive = sink;
224 c.cancelled().await;
225 cancelled.store(true, std::sync::atomic::Ordering::SeqCst);
226 });
227 Ok(rx)
228 }
229 }
230
231 let dir = tempdir().unwrap();
232 let header = SessionHeader {
233 agent: None,
234 session_id: "hung".into(),
235 created_at: chrono::Utc::now(),
236 cwd: dir.path().to_path_buf(),
237 parent_session_id: None,
238 contract_id: None,
239 work_item_id: None,
240 conversation: None,
241 contract: FrozenContract {
242 app_version: "0".into(),
243 provider: "hung".into(),
244 model: "m".into(),
245 route_ladder: Vec::new(),
246 route_objective: String::new(),
247 route_annotations: Vec::new(),
248 system_prompt: "sys".into(),
249 permission_mode: "full-access".into(),
250 capabilities: Vec::new(),
251 prompt_layers: Vec::new(),
252 },
253 };
254 let log = SessionLog::create(dir.path().join("s.jsonl"), header).unwrap();
255 let mut cfg = AgentConfig::new("sys");
256 cfg.max_retries = 0;
257 cfg.request_timeout = Some(std::time::Duration::from_millis(300));
258 cfg.run_retry_attempts = 0;
259 std::mem::forget(dir);
260 let mut agent = Agent::new(Arc::new(Hung(cancelled.clone())), log, cfg);
261
262 let start = Instant::now();
263 let outcome = agent
264 .run(
265 "go",
266 &Default::default(),
267 CancellationToken::new(),
268 mpsc::channel(64).0,
269 )
270 .await;
271 let elapsed = start.elapsed();
272
273 eprintln!("OUTCOME: {outcome:?}");
274 match outcome {
275 TurnOutcome::Failed { error } => {
276 assert!(error.to_string().contains("deadline"), "got: {error}");
277 }
278 other => panic!("expected deadline failure, got {other:?}"),
279 }
280 assert!(
281 elapsed < std::time::Duration::from_secs(3),
282 "watchdog must fire promptly, took {elapsed:?}"
283 );
284 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
285 assert!(cancelled.load(std::sync::atomic::Ordering::SeqCst));
286}
287
288/// Hangs (holds its stream open until cancelled) for the first
289/// `hangs_remaining` calls, then serves `good` (or a plain "done"). Counts
290/// every call and every attempt-cancellation it actually observed, so a
291/// test can tell a per-attempt child token from the run's own.
292struct HungThenGood {
293 hangs_remaining: Mutex<u32>,
294 attempts: Arc<std::sync::atomic::AtomicU32>,
295 attempt_cancellations: Arc<std::sync::atomic::AtomicU32>,
296 good: Mutex<Option<AssistantMessage>>,
297}
298
299#[async_trait::async_trait]
300impl Provider for HungThenGood {
301 fn name(&self) -> &str {
302 "hung"
303 }
304 async fn stream(&self, _r: ChatRequest, c: CancellationToken) -> Result<EventStream, LlmError> {
305 self.attempts
306 .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
307 let should_hang = {
308 let mut remaining = self.hangs_remaining.lock().unwrap();
309 if *remaining > 0 {
310 *remaining -= 1;
311 true
312 } else {
313 false
314 }
315 };
316 let (mut sink, rx) = stream::channel(8);
317 if should_hang {
318 sink.push(stream::StreamEvent::Start {
319 partial: AssistantMessage::empty("m"),
320 });
321 let attempt_cancellations = self.attempt_cancellations.clone();
322 tokio::spawn(async move {
323 let mut sink = sink;
324 c.cancelled().await;
325 attempt_cancellations.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
326 // A well-behaved provider always closes with a terminal
327 // event on cancellation (every real adapter does); a mock
328 // that just dropped the sink here would surface as
329 // "stream ended without a terminal event" instead of the
330 // Aborted this simulates.
331 sink.close_error(LlmError::Aborted { partial: None }).await;
332 });
333 } else {
334 let msg = self
335 .good
336 .lock()
337 .unwrap()
338 .take()
339 .unwrap_or_else(|| text_msg_for_hung("done"));
340 sink.push(stream::StreamEvent::Start {
341 partial: msg.clone(),
342 });
343 sink.close_message(msg).await;
344 }
345 Ok(rx)
346 }
347}
348
349fn text_msg_for_hung(t: &str) -> AssistantMessage {
350 AssistantMessage {
351 content: vec![ContentBlock::text(t)],
352 stop_reason: StopReason::EndTurn,
353 usage: Usage::default(),
354 model: "m".into(),
355 response_id: None,
356 }
357}
358
359fn hung_header(provider: &str, session_id: &str, dir: &std::path::Path) -> SessionHeader {
360 SessionHeader {
361 agent: None,
362 session_id: session_id.into(),
363 created_at: chrono::Utc::now(),
364 cwd: dir.to_path_buf(),
365 parent_session_id: None,
366 contract_id: None,
367 work_item_id: None,
368 conversation: None,
369 contract: FrozenContract {
370 app_version: "0".into(),
371 provider: provider.into(),
372 model: "m".into(),
373 route_ladder: Vec::new(),
374 route_objective: String::new(),
375 route_annotations: Vec::new(),
376 system_prompt: "sys".into(),
377 permission_mode: "full-access".into(),
378 capabilities: Vec::new(),
379 prompt_layers: Vec::new(),
380 },
381 }
382}
383
384/// A per-attempt timeout must never poison the run's own cancellation
385/// token: with retries enabled, every attempt independently hangs and
386/// times out, and the run still gets a full, bounded set of attempts
387/// (initial + `max_retries`) rather than the first timeout silently
388/// aborting the rest.
389#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
390async fn hung_provider_with_retries_gets_bounded_attempts_then_fails_with_deadline() {
391 let dir = tempdir().unwrap();
392 let log = SessionLog::create(
393 dir.path().join("s.jsonl"),
394 hung_header("hung", "hung-retries", dir.path()),
395 )
396 .unwrap();
397 let attempts = Arc::new(std::sync::atomic::AtomicU32::new(0));
398 let provider = Arc::new(HungThenGood {
399 hangs_remaining: Mutex::new(u32::MAX),
400 attempts: attempts.clone(),
401 attempt_cancellations: Arc::new(std::sync::atomic::AtomicU32::new(0)),
402 good: Mutex::new(None),
403 });
404 let mut cfg = AgentConfig::new("sys");
405 cfg.max_retries = 2;
406 cfg.retry_base_backoff_ms = 1;
407 cfg.request_timeout = Some(std::time::Duration::from_millis(150));
408 cfg.run_retry_attempts = 0;
409 std::mem::forget(dir);
410 let mut agent = Agent::new(provider, log, cfg);
411
412 let outcome = agent
413 .run(
414 "go",
415 &Default::default(),
416 CancellationToken::new(),
417 mpsc::channel(64).0,
418 )
419 .await;
420 match outcome {
421 TurnOutcome::Failed { error } => assert!(error.to_string().contains("deadline")),
422 other => panic!("expected deadline failure, got {other:?}"),
423 }
424 assert_eq!(
425 attempts.load(std::sync::atomic::Ordering::SeqCst),
426 3,
427 "initial attempt + 2 retries, each independently timing out"
428 );
429}
430
431/// A leg that always times out falls back to the next frozen-ladder leg
432/// instead of the run-level `cancel` token being poisoned by the first
433/// timeout (which used to make the very next retry/fallback wait observe
434/// an already-cancelled token and return Aborted instead of falling back).
435#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
436async fn a_hung_leg_falls_back_to_a_working_ladder_leg() {
437 let dir = tempdir().unwrap();
438 let log = SessionLog::create(
439 dir.path().join("s.jsonl"),
440 hung_header("hung", "hung-fallback", dir.path()),
441 )
442 .unwrap();
443 let hung_provider = Arc::new(HungThenGood {
444 hangs_remaining: Mutex::new(u32::MAX),
445 attempts: Arc::new(std::sync::atomic::AtomicU32::new(0)),
446 attempt_cancellations: Arc::new(std::sync::atomic::AtomicU32::new(0)),
447 good: Mutex::new(None),
448 });
449 let good_provider: Arc<dyn Provider> = Arc::new(Scripted {
450 steps: Mutex::new(VecDeque::from(vec![Step::Text("from the fallback".into())])),
451 });
452 let mut cfg = AgentConfig::new("sys");
453 cfg.max_retries = 0;
454 cfg.request_timeout = Some(std::time::Duration::from_millis(150));
455 cfg.run_retry_attempts = 0;
456 cfg.ladder = vec![(good_provider, "good-model".into())];
457 cfg.ladder_provider_names = vec!["good".into()];
458 std::mem::forget(dir);
459 let mut agent = Agent::new(hung_provider, log, cfg);
460
461 let outcome = agent
462 .run(
463 "go",
464 &Default::default(),
465 CancellationToken::new(),
466 mpsc::channel(64).0,
467 )
468 .await;
469 match outcome {
470 TurnOutcome::Completed { response } => {
471 assert_eq!(response.text_content(), "from the fallback");
472 }
473 other => panic!("expected the fallback leg to complete, got {other:?}"),
474 }
475}
476
477/// A step that exhausts its retry budget by timing out is a transient
478/// (`LlmError::Network`) failure at the run level too, since the run's own
479/// `cancel` was never touched by the per-attempt timeout: run-level
480/// endurance gets a real re-attempt instead of seeing an Aborted turn.
481#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
482async fn run_level_endurance_re_attempts_after_a_hung_step() {
483 let dir = tempdir().unwrap();
484 let log = SessionLog::create(
485 dir.path().join("s.jsonl"),
486 hung_header("hung", "hung-endurance", dir.path()),
487 )
488 .unwrap();
489 let attempts = Arc::new(std::sync::atomic::AtomicU32::new(0));
490 let provider = Arc::new(HungThenGood {
491 // Hangs on the first step-level attempt only; the run-level
492 // re-attempt's own first (and only) provider call succeeds.
493 hangs_remaining: Mutex::new(1),
494 attempts: attempts.clone(),
495 attempt_cancellations: Arc::new(std::sync::atomic::AtomicU32::new(0)),
496 good: Mutex::new(Some(text_msg_for_hung("recovered"))),
497 });
498 let mut cfg = AgentConfig::new("sys");
499 cfg.max_retries = 0;
500 cfg.request_timeout = Some(std::time::Duration::from_millis(150));
501 cfg.run_retry_attempts = 1;
502 cfg.run_retry_base_backoff_ms = 1;
503 std::mem::forget(dir);
504 let mut agent = Agent::new(provider, log, cfg);
505
506 let outcome = agent
507 .run(
508 "go",
509 &Default::default(),
510 CancellationToken::new(),
511 mpsc::channel(64).0,
512 )
513 .await;
514 match outcome {
515 TurnOutcome::Completed { response } => {
516 assert_eq!(response.text_content(), "recovered");
517 }
518 other => panic!("expected run-level endurance to recover, got {other:?}"),
519 }
520 assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 2);
521}
522
523/// A real user cancellation during a hung attempt still aborts immediately
524/// -- the per-attempt child token derived from `cancel` must observe the
525/// parent's cancellation (propagation is one-directional: parent cancels
526/// child, never the reverse).
527#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
528async fn user_cancel_during_a_hung_attempt_gives_aborted() {
529 let dir = tempdir().unwrap();
530 let log = SessionLog::create(
531 dir.path().join("s.jsonl"),
532 hung_header("hung", "hung-cancel", dir.path()),
533 )
534 .unwrap();
535 let provider = Arc::new(HungThenGood {
536 hangs_remaining: Mutex::new(u32::MAX),
537 attempts: Arc::new(std::sync::atomic::AtomicU32::new(0)),
538 attempt_cancellations: Arc::new(std::sync::atomic::AtomicU32::new(0)),
539 good: Mutex::new(None),
540 });
541 let mut cfg = AgentConfig::new("sys");
542 cfg.max_retries = 3;
543 // No request_timeout: only the user's own cancellation can end this.
544 cfg.run_retry_attempts = 0;
545 std::mem::forget(dir);
546 let mut agent = Agent::new(provider, log, cfg);
547
548 let cancel = CancellationToken::new();
549 let canceller = cancel.clone();
550 tokio::spawn(async move {
551 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
552 canceller.cancel();
553 });
554
555 let outcome = tokio::time::timeout(
556 std::time::Duration::from_secs(5),
557 agent.run("go", &Default::default(), cancel, mpsc::channel(64).0),
558 )
559 .await
560 .expect("a user cancel must not hang");
561 assert!(
562 matches!(outcome, TurnOutcome::Aborted { .. }),
563 "expected Aborted, got {outcome:?}"
564 );
565}
566

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, TurnOutcome};
13use vak_llm::stream;
14use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, StopReason, Usage};
15use vak_llm::{EventStream, LlmError, Provider};
16use vak_session::SessionLog;
17use vak_session::types::{FrozenContract, SessionHeader};
18
19enum Step {
20 Err(LlmError),
21 Text(String),
22}
23
24struct Scripted {
25 steps: Mutex<VecDeque<Step>>,
26}
27
28#[async_trait::async_trait]
29impl Provider for Scripted {
30 fn name(&self) -> &str {
31 "scripted"
32 }
33
34 async fn stream(
35 &self,
36 _request: ChatRequest,
37 _cancel: CancellationToken,
38 ) -> Result<EventStream, LlmError> {
39 let next = self.steps.lock().unwrap().pop_front();
40 let (mut sink, rx) = stream::channel(64);
41 match next {
42 Some(Step::Text(t)) => {
43 let msg = AssistantMessage {
44 content: vec![ContentBlock::text(t)],
45 stop_reason: StopReason::EndTurn,
46 usage: Usage::default(),
47 model: "test-model".into(),
48 response_id: None,
49 };
50 sink.push(stream::StreamEvent::Start {
51 partial: msg.clone(),
52 });
53 sink.close_message(msg).await;
54 }
55 Some(Step::Err(e)) => sink.close_error(e).await,
56 None => sink.close_error(LlmError::Parse("exhausted".into())).await,
57 }
58 Ok(rx)
59 }
60}
61
62fn build(
63 steps: Vec<Step>,
64 max_retries: u32,
65 base_ms: u64,
66 timeout: Option<std::time::Duration>,
67) -> Agent {
68 let dir = tempdir().unwrap();
69 let header = SessionHeader {
70 agent: None,
71 session_id: "rel".into(),
72 created_at: chrono::Utc::now(),
73 cwd: dir.path().to_path_buf(),
74 parent_session_id: None,
75 contract_id: None,
76 work_item_id: None,
77 conversation: None,
78 contract: FrozenContract {
79 app_version: "0".into(),
80 provider: "scripted".into(),
81 model: "test-model".into(),
82 route_ladder: Vec::new(),
83 route_objective: String::new(),
84 route_annotations: Vec::new(),
85 system_prompt: "sys".into(),
86 permission_mode: "full-access".into(),
87 capabilities: Vec::new(),
88 prompt_layers: Vec::new(),
89 },
90 };
91 let log = SessionLog::create(dir.path().join("s.jsonl"), header).unwrap();
92 let mut cfg = AgentConfig::new("sys");
93 cfg.max_retries = max_retries;
94 cfg.retry_base_backoff_ms = base_ms;
95 cfg.request_timeout = timeout;
96 // Step-level machinery is the system under test here; run-level
97 // endurance has its own tests (tests/run_endurance.rs).
98 cfg.run_retry_attempts = 0;
99 std::mem::forget(dir);
100 Agent::new(
101 Arc::new(Scripted {
102 steps: Mutex::new(steps.into_iter().collect()),
103 }),
104 log,
105 cfg,
106 )
107}
108
109#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
110async fn transient_errors_are_retried_until_success() {
111 let mut agent = build(
112 vec![
113 Step::Err(LlmError::RateLimit {
114 message: "slow down".into(),
115 retry_after_secs: None,
116 }),
117 Step::Err(LlmError::Overloaded("529".into())),
118 Step::Err(LlmError::Network("flap".into())),
119 Step::Text("finally".into()),
120 ],
121 3,
122 1, // near-zero backoff for the test
123 None,
124 );
125
126 let outcome = agent
127 .run(
128 "go",
129 &Default::default(),
130 CancellationToken::new(),
131 mpsc::channel(64).0,
132 )
133 .await;
134
135 match outcome {
136 TurnOutcome::Completed { response } => assert_eq!(response.text_content(), "finally"),
137 other => panic!("expected completed after retries, got {other:?}"),
138 }
139}
140
141#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
142async fn non_retryable_errors_fail_immediately() {
143 let mut agent = build(
144 vec![
145 Step::Err(LlmError::Auth("bad key".into())),
146 Step::Text("never".into()),
147 ],
148 3,
149 1,
150 None,
151 );
152
153 let start = Instant::now();
154 let outcome = agent
155 .run(
156 "go",
157 &Default::default(),
158 CancellationToken::new(),
159 mpsc::channel(64).0,
160 )
161 .await;
162 assert!(
163 start.elapsed() < std::time::Duration::from_millis(500),
164 "auth errors must not be retried"
165 );
166 match outcome {
167 TurnOutcome::Failed { error } => assert!(error.to_string().contains("bad key")),
168 other => panic!("expected failed, got {other:?}"),
169 }
170}
171
172#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
173async fn retry_budget_exhaustion_fails_with_last_error() {
174 let err_fn = || {
175 Step::Err(LlmError::RateLimit {
176 message: "still limited".into(),
177 retry_after_secs: None,
178 })
179 };
180 let mut agent = build(
181 vec![err_fn(), err_fn(), err_fn(), err_fn(), err_fn()],
182 2, // budget smaller than failure count
183 1,
184 None,
185 );
186
187 let outcome = agent
188 .run(
189 "go",
190 &Default::default(),
191 CancellationToken::new(),
192 mpsc::channel(64).0,
193 )
194 .await;
195 match outcome {
196 TurnOutcome::Failed { error } => assert!(error.to_string().contains("still limited")),
197 other => panic!("expected failed after budget, got {other:?}"),
198 }
199}
200
201#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
202async fn watchdog_deadline_converts_hung_step_into_retryable_failure() {
203 // A provider whose stream never completes: the deadline must fire.
204 let cancelled = Arc::new(std::sync::atomic::AtomicBool::new(false));
205 struct Hung(Arc<std::sync::atomic::AtomicBool>);
206 #[async_trait::async_trait]
207 impl Provider for Hung {
208 fn name(&self) -> &str {
209 "hung"
210 }
211 async fn stream(
212 &self,
213 _r: ChatRequest,
214 c: CancellationToken,
215 ) -> Result<EventStream, LlmError> {
216 let cancelled = self.0.clone();
217 let (mut sink, rx) = stream::channel(8);
218 sink.push(stream::StreamEvent::Start {
219 partial: AssistantMessage::empty("m"),
220 });
221 // Hold the sink open forever — simulates a stalled stream.
222 tokio::spawn(async move {
223 let _keep_alive = sink;
224 c.cancelled().await;
225 cancelled.store(true, std::sync::atomic::Ordering::SeqCst);
226 });
227 Ok(rx)
228 }
229 }
230
231 let dir = tempdir().unwrap();
232 let header = SessionHeader {
233 agent: None,
234 session_id: "hung".into(),
235 created_at: chrono::Utc::now(),
236 cwd: dir.path().to_path_buf(),
237 parent_session_id: None,
238 contract_id: None,
239 work_item_id: None,
240 conversation: None,
241 contract: FrozenContract {
242 app_version: "0".into(),
243 provider: "hung".into(),
244 model: "m".into(),
245 route_ladder: Vec::new(),
246 route_objective: String::new(),
247 route_annotations: Vec::new(),
248 system_prompt: "sys".into(),
249 permission_mode: "full-access".into(),
250 capabilities: Vec::new(),
251 prompt_layers: Vec::new(),
252 },
253 };
254 let log = SessionLog::create(dir.path().join("s.jsonl"), header).unwrap();
255 let mut cfg = AgentConfig::new("sys");
256 cfg.max_retries = 0;
257 cfg.request_timeout = Some(std::time::Duration::from_millis(300));
258 cfg.run_retry_attempts = 0;
259 std::mem::forget(dir);
260 let mut agent = Agent::new(Arc::new(Hung(cancelled.clone())), log, cfg);
261
262 let start = Instant::now();
263 let outcome = agent
264 .run(
265 "go",
266 &Default::default(),
267 CancellationToken::new(),
268 mpsc::channel(64).0,
269 )
270 .await;
271 let elapsed = start.elapsed();
272
273 eprintln!("OUTCOME: {outcome:?}");
274 match outcome {
275 TurnOutcome::Failed { error } => {
276 assert!(error.to_string().contains("deadline"), "got: {error}");
277 }
278 other => panic!("expected deadline failure, got {other:?}"),
279 }
280 assert!(
281 elapsed < std::time::Duration::from_secs(3),
282 "watchdog must fire promptly, took {elapsed:?}"
283 );
284 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
285 assert!(cancelled.load(std::sync::atomic::Ordering::SeqCst));
286}
287
288/// Hangs (holds its stream open until cancelled) for the first
289/// `hangs_remaining` calls, then serves `good` (or a plain "done"). Counts
290/// every call and every attempt-cancellation it actually observed, so a
291/// test can tell a per-attempt child token from the run's own.
292struct HungThenGood {
293 hangs_remaining: Mutex<u32>,
294 attempts: Arc<std::sync::atomic::AtomicU32>,
295 attempt_cancellations: Arc<std::sync::atomic::AtomicU32>,
296 good: Mutex<Option<AssistantMessage>>,
297}
298
299#[async_trait::async_trait]
300impl Provider for HungThenGood {
301 fn name(&self) -> &str {
302 "hung"
303 }
304 async fn stream(&self, _r: ChatRequest, c: CancellationToken) -> Result<EventStream, LlmError> {
305 self.attempts
306 .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
307 let should_hang = {
308 let mut remaining = self.hangs_remaining.lock().unwrap();
309 if *remaining > 0 {
310 *remaining -= 1;
311 true
312 } else {
313 false
314 }
315 };
316 let (mut sink, rx) = stream::channel(8);
317 if should_hang {
318 sink.push(stream::StreamEvent::Start {
319 partial: AssistantMessage::empty("m"),
320 });
321 let attempt_cancellations = self.attempt_cancellations.clone();
322 tokio::spawn(async move {
323 let mut sink = sink;
324 c.cancelled().await;
325 attempt_cancellations.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
326 // A well-behaved provider always closes with a terminal
327 // event on cancellation (every real adapter does); a mock
328 // that just dropped the sink here would surface as
329 // "stream ended without a terminal event" instead of the
330 // Aborted this simulates.
331 sink.close_error(LlmError::Aborted { partial: None }).await;
332 });
333 } else {
334 let msg = self
335 .good
336 .lock()
337 .unwrap()
338 .take()
339 .unwrap_or_else(|| text_msg_for_hung("done"));
340 sink.push(stream::StreamEvent::Start {
341 partial: msg.clone(),
342 });
343 sink.close_message(msg).await;
344 }
345 Ok(rx)
346 }
347}
348
349fn text_msg_for_hung(t: &str) -> AssistantMessage {
350 AssistantMessage {
351 content: vec![ContentBlock::text(t)],
352 stop_reason: StopReason::EndTurn,
353 usage: Usage::default(),
354 model: "m".into(),
355 response_id: None,
356 }
357}
358
359fn hung_header(provider: &str, session_id: &str, dir: &std::path::Path) -> SessionHeader {
360 SessionHeader {
361 agent: None,
362 session_id: session_id.into(),
363 created_at: chrono::Utc::now(),
364 cwd: dir.to_path_buf(),
365 parent_session_id: None,
366 contract_id: None,
367 work_item_id: None,
368 conversation: None,
369 contract: FrozenContract {
370 app_version: "0".into(),
371 provider: provider.into(),
372 model: "m".into(),
373 route_ladder: Vec::new(),
374 route_objective: String::new(),
375 route_annotations: Vec::new(),
376 system_prompt: "sys".into(),
377 permission_mode: "full-access".into(),
378 capabilities: Vec::new(),
379 prompt_layers: Vec::new(),
380 },
381 }
382}
383
384/// A per-attempt timeout must never poison the run's own cancellation
385/// token: with retries enabled, every attempt independently hangs and
386/// times out, and the run still gets a full, bounded set of attempts
387/// (initial + `max_retries`) rather than the first timeout silently
388/// aborting the rest.
389#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
390async fn hung_provider_with_retries_gets_bounded_attempts_then_fails_with_deadline() {
391 let dir = tempdir().unwrap();
392 let log = SessionLog::create(
393 dir.path().join("s.jsonl"),
394 hung_header("hung", "hung-retries", dir.path()),
395 )
396 .unwrap();
397 let attempts = Arc::new(std::sync::atomic::AtomicU32::new(0));
398 let provider = Arc::new(HungThenGood {
399 hangs_remaining: Mutex::new(u32::MAX),
400 attempts: attempts.clone(),
401 attempt_cancellations: Arc::new(std::sync::atomic::AtomicU32::new(0)),
402 good: Mutex::new(None),
403 });
404 let mut cfg = AgentConfig::new("sys");
405 cfg.max_retries = 2;
406 cfg.retry_base_backoff_ms = 1;
407 cfg.request_timeout = Some(std::time::Duration::from_millis(150));
408 cfg.run_retry_attempts = 0;
409 std::mem::forget(dir);
410 let mut agent = Agent::new(provider, log, cfg);
411
412 let outcome = agent
413 .run(
414 "go",
415 &Default::default(),
416 CancellationToken::new(),
417 mpsc::channel(64).0,
418 )
419 .await;
420 match outcome {
421 TurnOutcome::Failed { error } => assert!(error.to_string().contains("deadline")),
422 other => panic!("expected deadline failure, got {other:?}"),
423 }
424 assert_eq!(
425 attempts.load(std::sync::atomic::Ordering::SeqCst),
426 3,
427 "initial attempt + 2 retries, each independently timing out"
428 );
429}
430
431/// A leg that always times out falls back to the next frozen-ladder leg
432/// instead of the run-level `cancel` token being poisoned by the first
433/// timeout (which used to make the very next retry/fallback wait observe
434/// an already-cancelled token and return Aborted instead of falling back).
435#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
436async fn a_hung_leg_falls_back_to_a_working_ladder_leg() {
437 let dir = tempdir().unwrap();
438 let log = SessionLog::create(
439 dir.path().join("s.jsonl"),
440 hung_header("hung", "hung-fallback", dir.path()),
441 )
442 .unwrap();
443 let hung_provider = Arc::new(HungThenGood {
444 hangs_remaining: Mutex::new(u32::MAX),
445 attempts: Arc::new(std::sync::atomic::AtomicU32::new(0)),
446 attempt_cancellations: Arc::new(std::sync::atomic::AtomicU32::new(0)),
447 good: Mutex::new(None),
448 });
449 let good_provider: Arc<dyn Provider> = Arc::new(Scripted {
450 steps: Mutex::new(VecDeque::from(vec![Step::Text("from the fallback".into())])),
451 });
452 let mut cfg = AgentConfig::new("sys");
453 cfg.max_retries = 0;
454 cfg.request_timeout = Some(std::time::Duration::from_millis(150));
455 cfg.run_retry_attempts = 0;
456 cfg.ladder = vec![(good_provider, "good-model".into())];
457 cfg.ladder_provider_names = vec!["good".into()];
458 std::mem::forget(dir);
459 let mut agent = Agent::new(hung_provider, log, cfg);
460
461 let outcome = agent
462 .run(
463 "go",
464 &Default::default(),
465 CancellationToken::new(),
466 mpsc::channel(64).0,
467 )
468 .await;
469 match outcome {
470 TurnOutcome::Completed { response } => {
471 assert_eq!(response.text_content(), "from the fallback");
472 }
473 other => panic!("expected the fallback leg to complete, got {other:?}"),
474 }
475}
476
477/// A step that exhausts its retry budget by timing out is a transient
478/// (`LlmError::Network`) failure at the run level too, since the run's own
479/// `cancel` was never touched by the per-attempt timeout: run-level
480/// endurance gets a real re-attempt instead of seeing an Aborted turn.
481#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
482async fn run_level_endurance_re_attempts_after_a_hung_step() {
483 let dir = tempdir().unwrap();
484 let log = SessionLog::create(
485 dir.path().join("s.jsonl"),
486 hung_header("hung", "hung-endurance", dir.path()),
487 )
488 .unwrap();
489 let attempts = Arc::new(std::sync::atomic::AtomicU32::new(0));
490 let provider = Arc::new(HungThenGood {
491 // Hangs on the first step-level attempt only; the run-level
492 // re-attempt's own first (and only) provider call succeeds.
493 hangs_remaining: Mutex::new(1),
494 attempts: attempts.clone(),
495 attempt_cancellations: Arc::new(std::sync::atomic::AtomicU32::new(0)),
496 good: Mutex::new(Some(text_msg_for_hung("recovered"))),
497 });
498 let mut cfg = AgentConfig::new("sys");
499 cfg.max_retries = 0;
500 cfg.request_timeout = Some(std::time::Duration::from_millis(150));
501 cfg.run_retry_attempts = 1;
502 cfg.run_retry_base_backoff_ms = 1;
503 std::mem::forget(dir);
504 let mut agent = Agent::new(provider, log, cfg);
505
506 let outcome = agent
507 .run(
508 "go",
509 &Default::default(),
510 CancellationToken::new(),
511 mpsc::channel(64).0,
512 )
513 .await;
514 match outcome {
515 TurnOutcome::Completed { response } => {
516 assert_eq!(response.text_content(), "recovered");
517 }
518 other => panic!("expected run-level endurance to recover, got {other:?}"),
519 }
520 assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 2);
521}
522
523/// A real user cancellation during a hung attempt still aborts immediately
524/// -- the per-attempt child token derived from `cancel` must observe the
525/// parent's cancellation (propagation is one-directional: parent cancels
526/// child, never the reverse).
527#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
528async fn user_cancel_during_a_hung_attempt_gives_aborted() {
529 let dir = tempdir().unwrap();
530 let log = SessionLog::create(
531 dir.path().join("s.jsonl"),
532 hung_header("hung", "hung-cancel", dir.path()),
533 )
534 .unwrap();
535 let provider = Arc::new(HungThenGood {
536 hangs_remaining: Mutex::new(u32::MAX),
537 attempts: Arc::new(std::sync::atomic::AtomicU32::new(0)),
538 attempt_cancellations: Arc::new(std::sync::atomic::AtomicU32::new(0)),
539 good: Mutex::new(None),
540 });
541 let mut cfg = AgentConfig::new("sys");
542 cfg.max_retries = 3;
543 // No request_timeout: only the user's own cancellation can end this.
544 cfg.run_retry_attempts = 0;
545 std::mem::forget(dir);
546 let mut agent = Agent::new(provider, log, cfg);
547
548 let cancel = CancellationToken::new();
549 let canceller = cancel.clone();
550 tokio::spawn(async move {
551 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
552 canceller.cancel();
553 });
554
555 let outcome = tokio::time::timeout(
556 std::time::Duration::from_secs(5),
557 agent.run("go", &Default::default(), cancel, mpsc::channel(64).0),
558 )
559 .await
560 .expect("a user cancel must not hang");
561 assert!(
562 matches!(outcome, TurnOutcome::Aborted { .. }),
563 "expected Aborted, got {outcome:?}"
564 );
565}
566

Reliability

Indexing the workspace…

Reliability