Vakyartha documentation

1#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
2
3//! Budget admission (docs/design/15-reliability.md): denial fails the run
4//! permanently with a typed budget message unless the approver accepts
5//! the one-time raise Ask; unattended (no approver) always aborts.
6
7use std::sync::Arc;
8use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU32, Ordering};
9
10use tempfile::tempdir;
11use tokio::sync::mpsc;
12use tokio_util::sync::CancellationToken;
13
14use vak_agent::{Agent, AgentConfig, Approver, SpendCheck, SpendGate, SteeringQueues, TurnOutcome};
15use vak_llm::stream;
16use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, StopReason, Usage};
17use vak_llm::{EventStream, LlmError, Provider};
18use vak_permission::{Mode, PermissionEngine};
19use vak_session::types::{FrozenContract, SessionHeader};
20use vak_session::{SessionLog, SessionPath};
21
22fn text_msg(t: &str) -> AssistantMessage {
23 AssistantMessage {
24 content: vec![ContentBlock::text(t)],
25 stop_reason: StopReason::EndTurn,
26 usage: Usage {
27 input_tokens: 5,
28 output_tokens: 3,
29 ..Default::default()
30 },
31 model: "test-model".into(),
32 response_id: None,
33 }
34}
35
36struct Success;
37
38#[async_trait::async_trait]
39impl Provider for Success {
40 fn name(&self) -> &str {
41 "ok"
42 }
43
44 async fn stream(
45 &self,
46 _request: ChatRequest,
47 _cancel: CancellationToken,
48 ) -> Result<EventStream, LlmError> {
49 let (mut sink, rx) = stream::channel(64);
50 let m = text_msg("done");
51 sink.push(stream::StreamEvent::Start { partial: m.clone() });
52 sink.close_message(m).await;
53 Ok(rx)
54 }
55}
56
57/// Denies the first authorize call N times.
58struct FlakyGate {
59 remaining_denials: AtomicU8,
60}
61
62#[async_trait::async_trait]
63impl SpendGate for FlakyGate {
64 async fn authorize(&self, _check: &SpendCheck<'_>) -> Result<(), String> {
65 if self.remaining_denials.fetch_sub(1, Ordering::SeqCst) > 0 {
66 Err("run budget $1.00 would be exceeded".into())
67 } else {
68 Ok(())
69 }
70 }
71
72 fn record_settled(&self, _provider: &str, _model: &str, _session_id: &str, _usage: &Usage) {}
73}
74
75struct AutoApproveBudget;
76
77#[async_trait::async_trait]
78impl Approver for AutoApproveBudget {
79 async fn approve(&self, tool: &str, _args_json: &str, _reason: &str) -> bool {
80 tool == "finops-budget"
81 }
82}
83
84fn setup(
85 approver: Option<Arc<dyn Approver>>,
86 gate: Arc<dyn SpendGate>,
87) -> (Agent, tempfile::TempDir) {
88 let dir = tempdir().unwrap();
89 let cwd = dir.path().to_path_buf();
90 let header = SessionHeader {
91 agent: None,
92 session_id: "spend".into(),
93 created_at: chrono::Utc::now(),
94 cwd: cwd.clone(),
95 parent_session_id: None,
96 contract_id: None,
97 work_item_id: None,
98 conversation: None,
99 contract: FrozenContract {
100 app_version: "0".into(),
101 provider: "ok".into(),
102 model: "test-model".into(),
103 route_ladder: Vec::new(),
104 route_objective: String::new(),
105 route_annotations: Vec::new(),
106 system_prompt: "sys".into(),
107 permission_mode: "workspace-write".into(),
108 capabilities: Vec::new(),
109 prompt_layers: Vec::new(),
110 },
111 };
112 let home = cwd.join(".vak-home");
113 std::fs::create_dir_all(&home).unwrap();
114 let log =
115 SessionLog::create(SessionPath::new_session_file(&home, &cwd, "spend"), header).unwrap();
116 let mut cfg = AgentConfig::new("sys");
117 cfg.model = "test-model".into();
118 cfg.mode = Mode::FullAccess;
119 cfg.permission = Some(Arc::new(PermissionEngine::default()));
120 cfg.approver = approver;
121 cfg.spend_gate = Some(gate);
122 cfg.retry_base_backoff_ms = 1;
123 cfg.run_retry_base_backoff_ms = 1;
124 (Agent::new(Arc::new(Success), log, cfg), dir)
125}
126
127async fn run(agent: &mut Agent) -> TurnOutcome {
128 let (ev_tx, mut ev_rx) = mpsc::channel(256);
129 tokio::spawn(async move { while ev_rx.recv().await.is_some() {} });
130 let cancel = CancellationToken::new();
131 let steering = SteeringQueues::new();
132 agent.run("hello", &steering, cancel, ev_tx).await
133}
134
135#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
136async fn budget_denial_without_approver_fails_permanently() {
137 let (mut agent, _dir) = setup(
138 None,
139 Arc::new(FlakyGate {
140 remaining_denials: AtomicU8::new(9),
141 }),
142 );
143 match run(&mut agent).await {
144 TurnOutcome::Failed { error } => {
145 assert!(
146 error.to_string().contains("budget admission denied"),
147 "{error}"
148 );
149 }
150 other => panic!("expected Failed, got {other:?}"),
151 }
152}
153
154#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
155async fn approved_budget_ask_lets_the_run_proceed() {
156 let (mut agent, _dir) = setup(
157 Some(Arc::new(AutoApproveBudget)),
158 Arc::new(FlakyGate {
159 remaining_denials: AtomicU8::new(1),
160 }),
161 );
162 let outcome = run(&mut agent).await;
163 assert!(
164 matches!(outcome, TurnOutcome::Completed { .. }),
165 "approved ask must proceed, got {outcome:?}"
166 );
167}
168
169/// Always denies until on_budget_approved() flips it.
170struct CapGate {
171 raised: AtomicBool,
172 denials: AtomicU32,
173}
174
175#[async_trait::async_trait]
176impl SpendGate for CapGate {
177 async fn authorize(&self, _check: &SpendCheck<'_>) -> Result<(), String> {
178 if !self.raised.load(Ordering::SeqCst) {
179 self.denials.fetch_add(1, Ordering::SeqCst);
180 Err("run budget $1.00 would be exceeded".into())
181 } else {
182 Ok(())
183 }
184 }
185
186 fn on_budget_approved(&self) {
187 self.raised.store(true, Ordering::SeqCst);
188 }
189
190 fn record_settled(&self, _provider: &str, _model: &str, _session_id: &str, _usage: &Usage) {}
191}
192
193struct ApproveOnce {
194 remaining: AtomicU32,
195}
196
197#[async_trait::async_trait]
198impl Approver for ApproveOnce {
199 async fn approve(&self, tool: &str, _args_json: &str, _reason: &str) -> bool {
200 tool == "finops-budget" && self.remaining.fetch_sub(1, Ordering::SeqCst) > 0
201 }
202}
203
204/// Approval raises the cap for the WHOLE run: later dispatches pass
205/// without re-asking (raise-cap-once, docs/design/15-reliability.md).
206#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
207async fn budget_approval_raises_cap_for_rest_of_run() {
208 let gate = Arc::new(CapGate {
209 raised: AtomicBool::new(false),
210 denials: AtomicU32::new(0),
211 });
212 let dir = tempdir().unwrap();
213 let cwd = dir.path().to_path_buf();
214 let header = SessionHeader {
215 agent: None,
216 session_id: "raise".into(),
217 created_at: chrono::Utc::now(),
218 cwd: cwd.clone(),
219 parent_session_id: None,
220 contract_id: None,
221 work_item_id: None,
222 conversation: None,
223 contract: FrozenContract {
224 app_version: "0".into(),
225 provider: "ok".into(),
226 model: "test-model".into(),
227 route_ladder: Vec::new(),
228 route_objective: String::new(),
229 route_annotations: Vec::new(),
230 system_prompt: "sys".into(),
231 permission_mode: "workspace-write".into(),
232 capabilities: Vec::new(),
233 prompt_layers: Vec::new(),
234 },
235 };
236 let home = cwd.join(".vak-home");
237 std::fs::create_dir_all(&home).unwrap();
238 let log =
239 SessionLog::create(SessionPath::new_session_file(&home, &cwd, "raise"), header).unwrap();
240 let mut cfg = AgentConfig::new("sys");
241 cfg.model = "test-model".into();
242 cfg.mode = Mode::FullAccess;
243 cfg.permission = Some(Arc::new(PermissionEngine::default()));
244 cfg.approver = Some(Arc::new(ApproveOnce {
245 remaining: AtomicU32::new(1),
246 }));
247 cfg.spend_gate = Some(gate.clone() as Arc<dyn SpendGate>);
248 cfg.retry_base_backoff_ms = 1;
249 let mut agent = Agent::new(Arc::new(Success), log, cfg);
250
251 let (ev_tx, mut ev_rx) = mpsc::channel(256);
252 tokio::spawn(async move { while ev_rx.recv().await.is_some() {} });
253 let outcome = agent
254 .run(
255 "hello",
256 &SteeringQueues::new(),
257 CancellationToken::new(),
258 ev_tx,
259 )
260 .await;
261 assert!(matches!(outcome, TurnOutcome::Completed { .. }));
262 assert_eq!(gate.denials.load(Ordering::SeqCst), 1, "exactly one denial");
263}
264

1#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
2
3//! Budget admission (docs/design/15-reliability.md): denial fails the run
4//! permanently with a typed budget message unless the approver accepts
5//! the one-time raise Ask; unattended (no approver) always aborts.
6
7use std::sync::Arc;
8use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU32, Ordering};
9
10use tempfile::tempdir;
11use tokio::sync::mpsc;
12use tokio_util::sync::CancellationToken;
13
14use vak_agent::{Agent, AgentConfig, Approver, SpendCheck, SpendGate, SteeringQueues, TurnOutcome};
15use vak_llm::stream;
16use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, StopReason, Usage};
17use vak_llm::{EventStream, LlmError, Provider};
18use vak_permission::{Mode, PermissionEngine};
19use vak_session::types::{FrozenContract, SessionHeader};
20use vak_session::{SessionLog, SessionPath};
21
22fn text_msg(t: &str) -> AssistantMessage {
23 AssistantMessage {
24 content: vec![ContentBlock::text(t)],
25 stop_reason: StopReason::EndTurn,
26 usage: Usage {
27 input_tokens: 5,
28 output_tokens: 3,
29 ..Default::default()
30 },
31 model: "test-model".into(),
32 response_id: None,
33 }
34}
35
36struct Success;
37
38#[async_trait::async_trait]
39impl Provider for Success {
40 fn name(&self) -> &str {
41 "ok"
42 }
43
44 async fn stream(
45 &self,
46 _request: ChatRequest,
47 _cancel: CancellationToken,
48 ) -> Result<EventStream, LlmError> {
49 let (mut sink, rx) = stream::channel(64);
50 let m = text_msg("done");
51 sink.push(stream::StreamEvent::Start { partial: m.clone() });
52 sink.close_message(m).await;
53 Ok(rx)
54 }
55}
56
57/// Denies the first authorize call N times.
58struct FlakyGate {
59 remaining_denials: AtomicU8,
60}
61
62#[async_trait::async_trait]
63impl SpendGate for FlakyGate {
64 async fn authorize(&self, _check: &SpendCheck<'_>) -> Result<(), String> {
65 if self.remaining_denials.fetch_sub(1, Ordering::SeqCst) > 0 {
66 Err("run budget $1.00 would be exceeded".into())
67 } else {
68 Ok(())
69 }
70 }
71
72 fn record_settled(&self, _provider: &str, _model: &str, _session_id: &str, _usage: &Usage) {}
73}
74
75struct AutoApproveBudget;
76
77#[async_trait::async_trait]
78impl Approver for AutoApproveBudget {
79 async fn approve(&self, tool: &str, _args_json: &str, _reason: &str) -> bool {
80 tool == "finops-budget"
81 }
82}
83
84fn setup(
85 approver: Option<Arc<dyn Approver>>,
86 gate: Arc<dyn SpendGate>,
87) -> (Agent, tempfile::TempDir) {
88 let dir = tempdir().unwrap();
89 let cwd = dir.path().to_path_buf();
90 let header = SessionHeader {
91 agent: None,
92 session_id: "spend".into(),
93 created_at: chrono::Utc::now(),
94 cwd: cwd.clone(),
95 parent_session_id: None,
96 contract_id: None,
97 work_item_id: None,
98 conversation: None,
99 contract: FrozenContract {
100 app_version: "0".into(),
101 provider: "ok".into(),
102 model: "test-model".into(),
103 route_ladder: Vec::new(),
104 route_objective: String::new(),
105 route_annotations: Vec::new(),
106 system_prompt: "sys".into(),
107 permission_mode: "workspace-write".into(),
108 capabilities: Vec::new(),
109 prompt_layers: Vec::new(),
110 },
111 };
112 let home = cwd.join(".vak-home");
113 std::fs::create_dir_all(&home).unwrap();
114 let log =
115 SessionLog::create(SessionPath::new_session_file(&home, &cwd, "spend"), header).unwrap();
116 let mut cfg = AgentConfig::new("sys");
117 cfg.model = "test-model".into();
118 cfg.mode = Mode::FullAccess;
119 cfg.permission = Some(Arc::new(PermissionEngine::default()));
120 cfg.approver = approver;
121 cfg.spend_gate = Some(gate);
122 cfg.retry_base_backoff_ms = 1;
123 cfg.run_retry_base_backoff_ms = 1;
124 (Agent::new(Arc::new(Success), log, cfg), dir)
125}
126
127async fn run(agent: &mut Agent) -> TurnOutcome {
128 let (ev_tx, mut ev_rx) = mpsc::channel(256);
129 tokio::spawn(async move { while ev_rx.recv().await.is_some() {} });
130 let cancel = CancellationToken::new();
131 let steering = SteeringQueues::new();
132 agent.run("hello", &steering, cancel, ev_tx).await
133}
134
135#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
136async fn budget_denial_without_approver_fails_permanently() {
137 let (mut agent, _dir) = setup(
138 None,
139 Arc::new(FlakyGate {
140 remaining_denials: AtomicU8::new(9),
141 }),
142 );
143 match run(&mut agent).await {
144 TurnOutcome::Failed { error } => {
145 assert!(
146 error.to_string().contains("budget admission denied"),
147 "{error}"
148 );
149 }
150 other => panic!("expected Failed, got {other:?}"),
151 }
152}
153
154#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
155async fn approved_budget_ask_lets_the_run_proceed() {
156 let (mut agent, _dir) = setup(
157 Some(Arc::new(AutoApproveBudget)),
158 Arc::new(FlakyGate {
159 remaining_denials: AtomicU8::new(1),
160 }),
161 );
162 let outcome = run(&mut agent).await;
163 assert!(
164 matches!(outcome, TurnOutcome::Completed { .. }),
165 "approved ask must proceed, got {outcome:?}"
166 );
167}
168
169/// Always denies until on_budget_approved() flips it.
170struct CapGate {
171 raised: AtomicBool,
172 denials: AtomicU32,
173}
174
175#[async_trait::async_trait]
176impl SpendGate for CapGate {
177 async fn authorize(&self, _check: &SpendCheck<'_>) -> Result<(), String> {
178 if !self.raised.load(Ordering::SeqCst) {
179 self.denials.fetch_add(1, Ordering::SeqCst);
180 Err("run budget $1.00 would be exceeded".into())
181 } else {
182 Ok(())
183 }
184 }
185
186 fn on_budget_approved(&self) {
187 self.raised.store(true, Ordering::SeqCst);
188 }
189
190 fn record_settled(&self, _provider: &str, _model: &str, _session_id: &str, _usage: &Usage) {}
191}
192
193struct ApproveOnce {
194 remaining: AtomicU32,
195}
196
197#[async_trait::async_trait]
198impl Approver for ApproveOnce {
199 async fn approve(&self, tool: &str, _args_json: &str, _reason: &str) -> bool {
200 tool == "finops-budget" && self.remaining.fetch_sub(1, Ordering::SeqCst) > 0
201 }
202}
203
204/// Approval raises the cap for the WHOLE run: later dispatches pass
205/// without re-asking (raise-cap-once, docs/design/15-reliability.md).
206#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
207async fn budget_approval_raises_cap_for_rest_of_run() {
208 let gate = Arc::new(CapGate {
209 raised: AtomicBool::new(false),
210 denials: AtomicU32::new(0),
211 });
212 let dir = tempdir().unwrap();
213 let cwd = dir.path().to_path_buf();
214 let header = SessionHeader {
215 agent: None,
216 session_id: "raise".into(),
217 created_at: chrono::Utc::now(),
218 cwd: cwd.clone(),
219 parent_session_id: None,
220 contract_id: None,
221 work_item_id: None,
222 conversation: None,
223 contract: FrozenContract {
224 app_version: "0".into(),
225 provider: "ok".into(),
226 model: "test-model".into(),
227 route_ladder: Vec::new(),
228 route_objective: String::new(),
229 route_annotations: Vec::new(),
230 system_prompt: "sys".into(),
231 permission_mode: "workspace-write".into(),
232 capabilities: Vec::new(),
233 prompt_layers: Vec::new(),
234 },
235 };
236 let home = cwd.join(".vak-home");
237 std::fs::create_dir_all(&home).unwrap();
238 let log =
239 SessionLog::create(SessionPath::new_session_file(&home, &cwd, "raise"), header).unwrap();
240 let mut cfg = AgentConfig::new("sys");
241 cfg.model = "test-model".into();
242 cfg.mode = Mode::FullAccess;
243 cfg.permission = Some(Arc::new(PermissionEngine::default()));
244 cfg.approver = Some(Arc::new(ApproveOnce {
245 remaining: AtomicU32::new(1),
246 }));
247 cfg.spend_gate = Some(gate.clone() as Arc<dyn SpendGate>);
248 cfg.retry_base_backoff_ms = 1;
249 let mut agent = Agent::new(Arc::new(Success), log, cfg);
250
251 let (ev_tx, mut ev_rx) = mpsc::channel(256);
252 tokio::spawn(async move { while ev_rx.recv().await.is_some() {} });
253 let outcome = agent
254 .run(
255 "hello",
256 &SteeringQueues::new(),
257 CancellationToken::new(),
258 ev_tx,
259 )
260 .await;
261 assert!(matches!(outcome, TurnOutcome::Completed { .. }));
262 assert_eq!(gate.denials.load(Ordering::SeqCst), 1, "exactly one denial");
263}
264

Spend Gate

Indexing the workspace…

Spend Gate