Vakyartha documentation

1#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
2
3use std::sync::{Arc, Mutex};
4
5use tokio::sync::mpsc;
6use tokio_util::sync::CancellationToken;
7
8use tempfile::tempdir;
9
10use vak_agent::{Agent, AgentConfig, CircuitBreaker, CircuitBreakerConfig, TurnOutcome};
11use vak_llm::stream;
12use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, StopReason, Usage};
13use vak_llm::{EventStream, LlmError, Provider};
14use vak_session::SessionLog;
15use vak_session::types::{FrozenContract, SessionHeader};
16
17#[test]
18fn breaker_opens_after_threshold_and_half_closes_after_cooldown() {
19 let br = CircuitBreaker::new(CircuitBreakerConfig {
20 threshold: 3,
21 cooldown: std::time::Duration::from_millis(120),
22 });
23
24 assert!(br.check().is_ok());
25 br.record_failure();
26 br.record_failure();
27 assert!(br.check().is_ok(), "below threshold");
28
29 br.record_failure();
30 assert!(br.check().is_err(), "open at threshold");
31
32 std::thread::sleep(std::time::Duration::from_millis(150));
33 assert!(br.check().is_ok(), "cooldown half-closes");
34
35 // A success resets the failure count.
36 br.record_success();
37 br.record_failure();
38 br.record_failure();
39 assert!(br.check().is_ok(), "success reset the counter");
40}
41
42#[test]
43fn half_open_breaker_allows_only_one_probe() {
44 let br = CircuitBreaker::new(CircuitBreakerConfig {
45 threshold: 1,
46 cooldown: std::time::Duration::from_millis(1),
47 });
48 br.record_failure();
49 std::thread::sleep(std::time::Duration::from_millis(5));
50 assert!(br.check().is_ok());
51 assert!(
52 br.check().is_err(),
53 "a second caller must not race the probe"
54 );
55 br.record_success();
56 assert!(br.check().is_ok());
57}
58
59#[test]
60fn non_retryable_failures_do_not_trip_the_breaker() {
61 let br = CircuitBreaker::new(CircuitBreakerConfig {
62 threshold: 2,
63 cooldown: std::time::Duration::from_secs(60),
64 });
65 // Simulate auth failures: they never call record_failure.
66 for _ in 0..10 {
67 assert!(br.check().is_ok());
68 }
69 assert!(br.check().is_ok());
70}
71
72#[test]
73fn provider_circuits_do_not_poison_healthy_fallbacks() {
74 let br = CircuitBreaker::new(CircuitBreakerConfig {
75 threshold: 2,
76 cooldown: std::time::Duration::from_secs(60),
77 });
78 br.record_failure_key("ollama");
79 br.record_failure_key("ollama");
80 assert!(br.check_key("ollama").is_err());
81 assert!(br.check_key("anthropic").is_ok());
82 br.record_success_key("anthropic");
83 assert!(br.check_key("anthropic").is_ok());
84}
85
86#[test]
87fn credential_circuits_are_independent_within_one_provider() {
88 let br = CircuitBreaker::new(CircuitBreakerConfig {
89 threshold: 1,
90 cooldown: std::time::Duration::from_secs(60),
91 });
92 br.record_failure_key("openai:key-a");
93 assert!(br.check_key("openai:key-a").is_err());
94 assert!(br.check_key("openai:key-b").is_ok());
95}
96
97struct Scripted {
98 calls: Arc<Mutex<u32>>,
99 fail_with: LlmError,
100}
101
102#[async_trait::async_trait]
103impl Provider for Scripted {
104 fn name(&self) -> &str {
105 "scripted"
106 }
107
108 async fn stream(
109 &self,
110 _request: ChatRequest,
111 _cancel: CancellationToken,
112 ) -> Result<EventStream, LlmError> {
113 *self.calls.lock().unwrap() += 1;
114 let (mut sink, rx) = stream::channel(8);
115 sink.close_error(self.fail_with.clone()).await;
116 Ok(rx)
117 }
118}
119
120fn text_msg(t: &str) -> AssistantMessage {
121 AssistantMessage {
122 content: vec![ContentBlock::text(t)],
123 stop_reason: StopReason::EndTurn,
124 usage: Usage::default(),
125 model: "test-model".into(),
126 response_id: None,
127 }
128}
129
130fn build_agent(
131 provider: Arc<dyn Provider>,
132 breaker: Arc<CircuitBreaker>,
133 session_id: &str,
134) -> Agent {
135 let dir = tempdir().unwrap();
136 let header = SessionHeader {
137 agent: None,
138 session_id: session_id.into(),
139 created_at: chrono::Utc::now(),
140 cwd: dir.path().to_path_buf(),
141 parent_session_id: None,
142 contract_id: None,
143 work_item_id: None,
144 conversation: None,
145 contract: FrozenContract {
146 app_version: "0".into(),
147 provider: "scripted".into(),
148 model: "test-model".into(),
149 route_ladder: Vec::new(),
150 route_objective: String::new(),
151 route_annotations: Vec::new(),
152 system_prompt: "sys".into(),
153 permission_mode: "full-access".into(),
154 capabilities: Vec::new(),
155 prompt_layers: Vec::new(),
156 },
157 };
158 let log = SessionLog::create(dir.path().join("s.jsonl"), header).unwrap();
159 let mut cfg = AgentConfig::new("sys");
160 cfg.max_retries = 1; // one retry per run => 2 failures per run
161 cfg.retry_base_backoff_ms = 1;
162 // Fail-fast contract under test: endurance disabled so step exhaustion
163 // ends the run immediately (endurance has its own test file).
164 cfg.run_retry_attempts = 0;
165 cfg.circuit_breaker = Some(breaker);
166 std::mem::forget(dir);
167 Agent::new(provider, log, cfg)
168}
169
170#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
171async fn second_run_fails_fast_when_circuit_is_open() {
172 let calls = Arc::new(Mutex::new(0u32));
173 let provider = Arc::new(Scripted {
174 calls: calls.clone(),
175 // Blind network loss: the only class that trips the breaker.
176 fail_with: LlmError::Network("provider unreachable".into()),
177 });
178 let breaker = Arc::new(CircuitBreaker::new(CircuitBreakerConfig {
179 threshold: 4,
180 cooldown: std::time::Duration::from_millis(200),
181 }));
182
183 // Run 1: 2 attempts (1 + 1 retry) => 2 calls, breaker counts 2 failures.
184 let mut a1 = build_agent(provider.clone(), breaker.clone(), "run-1");
185 let outcome = a1
186 .run(
187 "go",
188 &Default::default(),
189 CancellationToken::new(),
190 mpsc::channel(64).0,
191 )
192 .await;
193 assert!(matches!(outcome, TurnOutcome::Failed { .. }));
194 assert_eq!(*calls.lock().unwrap(), 2);
195
196 // Run 2 (fresh agent, SAME breaker): 2 more failures => 4 total >= 4 → open.
197 let mut a2 = build_agent(provider.clone(), breaker.clone(), "run-2");
198 let outcome = a2
199 .run(
200 "go",
201 &Default::default(),
202 CancellationToken::new(),
203 mpsc::channel(64).0,
204 )
205 .await;
206 assert!(matches!(outcome, TurnOutcome::Failed { .. }));
207 assert_eq!(*calls.lock().unwrap(), 4);
208
209 // Run 3: circuit is open — fails fast WITHOUT touching the provider.
210 let mut a3 = build_agent(provider.clone(), breaker.clone(), "run-3");
211 let outcome = a3
212 .run(
213 "go",
214 &Default::default(),
215 CancellationToken::new(),
216 mpsc::channel(64).0,
217 )
218 .await;
219 match outcome {
220 TurnOutcome::Failed { error } => {
221 assert!(error.to_string().contains("circuit open"), "got: {error}");
222 }
223 other => panic!("expected fast failure, got {other:?}"),
224 }
225 assert_eq!(
226 *calls.lock().unwrap(),
227 4,
228 "open circuit must not reach the provider"
229 );
230
231 // A healthy provider sharing the breaker heals it on success.
232 struct Healthy;
233 #[async_trait::async_trait]
234 impl Provider for Healthy {
235 fn name(&self) -> &str {
236 "healthy"
237 }
238 async fn stream(
239 &self,
240 _r: ChatRequest,
241 _c: CancellationToken,
242 ) -> Result<EventStream, LlmError> {
243 let (mut sink, rx) = stream::channel(8);
244 let msg = text_msg("recovered");
245 sink.push(stream::StreamEvent::Start {
246 partial: msg.clone(),
247 });
248 sink.close_message(msg).await;
249 Ok(rx)
250 }
251 }
252 // After the cooldown the circuit half-closes: one probe gets through,
253 // and a healthy provider heals it on success.
254 std::thread::sleep(std::time::Duration::from_millis(250));
255 let mut a4 = build_agent(Arc::new(Healthy), breaker.clone(), "run-4");
256 let outcome = a4
257 .run(
258 "go",
259 &Default::default(),
260 CancellationToken::new(),
261 mpsc::channel(64).0,
262 )
263 .await;
264 assert!(matches!(outcome, TurnOutcome::Completed { .. }));
265 assert!(breaker.check().is_ok(), "success must close the circuit");
266}
267

1#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
2
3use std::sync::{Arc, Mutex};
4
5use tokio::sync::mpsc;
6use tokio_util::sync::CancellationToken;
7
8use tempfile::tempdir;
9
10use vak_agent::{Agent, AgentConfig, CircuitBreaker, CircuitBreakerConfig, TurnOutcome};
11use vak_llm::stream;
12use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, StopReason, Usage};
13use vak_llm::{EventStream, LlmError, Provider};
14use vak_session::SessionLog;
15use vak_session::types::{FrozenContract, SessionHeader};
16
17#[test]
18fn breaker_opens_after_threshold_and_half_closes_after_cooldown() {
19 let br = CircuitBreaker::new(CircuitBreakerConfig {
20 threshold: 3,
21 cooldown: std::time::Duration::from_millis(120),
22 });
23
24 assert!(br.check().is_ok());
25 br.record_failure();
26 br.record_failure();
27 assert!(br.check().is_ok(), "below threshold");
28
29 br.record_failure();
30 assert!(br.check().is_err(), "open at threshold");
31
32 std::thread::sleep(std::time::Duration::from_millis(150));
33 assert!(br.check().is_ok(), "cooldown half-closes");
34
35 // A success resets the failure count.
36 br.record_success();
37 br.record_failure();
38 br.record_failure();
39 assert!(br.check().is_ok(), "success reset the counter");
40}
41
42#[test]
43fn half_open_breaker_allows_only_one_probe() {
44 let br = CircuitBreaker::new(CircuitBreakerConfig {
45 threshold: 1,
46 cooldown: std::time::Duration::from_millis(1),
47 });
48 br.record_failure();
49 std::thread::sleep(std::time::Duration::from_millis(5));
50 assert!(br.check().is_ok());
51 assert!(
52 br.check().is_err(),
53 "a second caller must not race the probe"
54 );
55 br.record_success();
56 assert!(br.check().is_ok());
57}
58
59#[test]
60fn non_retryable_failures_do_not_trip_the_breaker() {
61 let br = CircuitBreaker::new(CircuitBreakerConfig {
62 threshold: 2,
63 cooldown: std::time::Duration::from_secs(60),
64 });
65 // Simulate auth failures: they never call record_failure.
66 for _ in 0..10 {
67 assert!(br.check().is_ok());
68 }
69 assert!(br.check().is_ok());
70}
71
72#[test]
73fn provider_circuits_do_not_poison_healthy_fallbacks() {
74 let br = CircuitBreaker::new(CircuitBreakerConfig {
75 threshold: 2,
76 cooldown: std::time::Duration::from_secs(60),
77 });
78 br.record_failure_key("ollama");
79 br.record_failure_key("ollama");
80 assert!(br.check_key("ollama").is_err());
81 assert!(br.check_key("anthropic").is_ok());
82 br.record_success_key("anthropic");
83 assert!(br.check_key("anthropic").is_ok());
84}
85
86#[test]
87fn credential_circuits_are_independent_within_one_provider() {
88 let br = CircuitBreaker::new(CircuitBreakerConfig {
89 threshold: 1,
90 cooldown: std::time::Duration::from_secs(60),
91 });
92 br.record_failure_key("openai:key-a");
93 assert!(br.check_key("openai:key-a").is_err());
94 assert!(br.check_key("openai:key-b").is_ok());
95}
96
97struct Scripted {
98 calls: Arc<Mutex<u32>>,
99 fail_with: LlmError,
100}
101
102#[async_trait::async_trait]
103impl Provider for Scripted {
104 fn name(&self) -> &str {
105 "scripted"
106 }
107
108 async fn stream(
109 &self,
110 _request: ChatRequest,
111 _cancel: CancellationToken,
112 ) -> Result<EventStream, LlmError> {
113 *self.calls.lock().unwrap() += 1;
114 let (mut sink, rx) = stream::channel(8);
115 sink.close_error(self.fail_with.clone()).await;
116 Ok(rx)
117 }
118}
119
120fn text_msg(t: &str) -> AssistantMessage {
121 AssistantMessage {
122 content: vec![ContentBlock::text(t)],
123 stop_reason: StopReason::EndTurn,
124 usage: Usage::default(),
125 model: "test-model".into(),
126 response_id: None,
127 }
128}
129
130fn build_agent(
131 provider: Arc<dyn Provider>,
132 breaker: Arc<CircuitBreaker>,
133 session_id: &str,
134) -> Agent {
135 let dir = tempdir().unwrap();
136 let header = SessionHeader {
137 agent: None,
138 session_id: session_id.into(),
139 created_at: chrono::Utc::now(),
140 cwd: dir.path().to_path_buf(),
141 parent_session_id: None,
142 contract_id: None,
143 work_item_id: None,
144 conversation: None,
145 contract: FrozenContract {
146 app_version: "0".into(),
147 provider: "scripted".into(),
148 model: "test-model".into(),
149 route_ladder: Vec::new(),
150 route_objective: String::new(),
151 route_annotations: Vec::new(),
152 system_prompt: "sys".into(),
153 permission_mode: "full-access".into(),
154 capabilities: Vec::new(),
155 prompt_layers: Vec::new(),
156 },
157 };
158 let log = SessionLog::create(dir.path().join("s.jsonl"), header).unwrap();
159 let mut cfg = AgentConfig::new("sys");
160 cfg.max_retries = 1; // one retry per run => 2 failures per run
161 cfg.retry_base_backoff_ms = 1;
162 // Fail-fast contract under test: endurance disabled so step exhaustion
163 // ends the run immediately (endurance has its own test file).
164 cfg.run_retry_attempts = 0;
165 cfg.circuit_breaker = Some(breaker);
166 std::mem::forget(dir);
167 Agent::new(provider, log, cfg)
168}
169
170#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
171async fn second_run_fails_fast_when_circuit_is_open() {
172 let calls = Arc::new(Mutex::new(0u32));
173 let provider = Arc::new(Scripted {
174 calls: calls.clone(),
175 // Blind network loss: the only class that trips the breaker.
176 fail_with: LlmError::Network("provider unreachable".into()),
177 });
178 let breaker = Arc::new(CircuitBreaker::new(CircuitBreakerConfig {
179 threshold: 4,
180 cooldown: std::time::Duration::from_millis(200),
181 }));
182
183 // Run 1: 2 attempts (1 + 1 retry) => 2 calls, breaker counts 2 failures.
184 let mut a1 = build_agent(provider.clone(), breaker.clone(), "run-1");
185 let outcome = a1
186 .run(
187 "go",
188 &Default::default(),
189 CancellationToken::new(),
190 mpsc::channel(64).0,
191 )
192 .await;
193 assert!(matches!(outcome, TurnOutcome::Failed { .. }));
194 assert_eq!(*calls.lock().unwrap(), 2);
195
196 // Run 2 (fresh agent, SAME breaker): 2 more failures => 4 total >= 4 → open.
197 let mut a2 = build_agent(provider.clone(), breaker.clone(), "run-2");
198 let outcome = a2
199 .run(
200 "go",
201 &Default::default(),
202 CancellationToken::new(),
203 mpsc::channel(64).0,
204 )
205 .await;
206 assert!(matches!(outcome, TurnOutcome::Failed { .. }));
207 assert_eq!(*calls.lock().unwrap(), 4);
208
209 // Run 3: circuit is open — fails fast WITHOUT touching the provider.
210 let mut a3 = build_agent(provider.clone(), breaker.clone(), "run-3");
211 let outcome = a3
212 .run(
213 "go",
214 &Default::default(),
215 CancellationToken::new(),
216 mpsc::channel(64).0,
217 )
218 .await;
219 match outcome {
220 TurnOutcome::Failed { error } => {
221 assert!(error.to_string().contains("circuit open"), "got: {error}");
222 }
223 other => panic!("expected fast failure, got {other:?}"),
224 }
225 assert_eq!(
226 *calls.lock().unwrap(),
227 4,
228 "open circuit must not reach the provider"
229 );
230
231 // A healthy provider sharing the breaker heals it on success.
232 struct Healthy;
233 #[async_trait::async_trait]
234 impl Provider for Healthy {
235 fn name(&self) -> &str {
236 "healthy"
237 }
238 async fn stream(
239 &self,
240 _r: ChatRequest,
241 _c: CancellationToken,
242 ) -> Result<EventStream, LlmError> {
243 let (mut sink, rx) = stream::channel(8);
244 let msg = text_msg("recovered");
245 sink.push(stream::StreamEvent::Start {
246 partial: msg.clone(),
247 });
248 sink.close_message(msg).await;
249 Ok(rx)
250 }
251 }
252 // After the cooldown the circuit half-closes: one probe gets through,
253 // and a healthy provider heals it on success.
254 std::thread::sleep(std::time::Duration::from_millis(250));
255 let mut a4 = build_agent(Arc::new(Healthy), breaker.clone(), "run-4");
256 let outcome = a4
257 .run(
258 "go",
259 &Default::default(),
260 CancellationToken::new(),
261 mpsc::channel(64).0,
262 )
263 .await;
264 assert!(matches!(outcome, TurnOutcome::Completed { .. }));
265 assert!(breaker.check().is_ok(), "success must close the circuit");
266}
267

Indexing the workspace…