1. 1#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
2. 2
3. 3use std::collections::VecDeque;
4. 4use std::sync::{Arc, Mutex};
5. 5
6. 6use tokio::sync::mpsc;
7. 7use tokio_util::sync::CancellationToken;
8. 8
9. 9use tempfile::tempdir;
10. 10
11. 11use vak_agent::{Agent, AgentConfig, TurnOutcome};
12. 12use vak_hooks::HookDef;
13. 13use vak_llm::stream;
14. 14use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, StopReason, Usage};
15. 15use vak_llm::{EventStream, LlmError, Provider};
16. 16use vak_session::SessionLog;
17. 17use vak_session::types::{FrozenContract, SessionHeader};
18. 18use vak_tools::bash::BashTool;
19. 19
20. 20struct Scripted {
21. 21 responses: Mutex<VecDeque<AssistantMessage>>,
22. 22}
23. 23
24. 24#[async_trait::async_trait]
25. 25impl Provider for Scripted {
26. 26 fn name(&self) -> &str {
27. 27 "scripted"
28. 28 }
29. 29
30. 30 async fn stream(
31. 31 &self,
32. 32 _request: ChatRequest,
33. 33 _cancel: CancellationToken,
34. 34 ) -> Result<EventStream, LlmError> {
35. 35 let next = self.responses.lock().unwrap().pop_front();
36. 36 let (mut sink, rx) = stream::channel(64);
37. 37 match next {
38. 38 Some(m) => {
39. 39 sink.push(stream::StreamEvent::Start { partial: m.clone() });
40. 40 sink.close_message(m).await;
41. 41 }
42. 42 None => sink.close_error(LlmError::Parse("exhausted".into())).await,
43. 43 }
44. 44 Ok(rx)
45. 45 }
46. 46}
47. 47
48. 48fn text_msg(t: &str) -> AssistantMessage {
49. 49 AssistantMessage {
50. 50 content: vec![ContentBlock::text(t)],
51. 51 stop_reason: StopReason::EndTurn,
52. 52 usage: Usage {
53. 53 input_tokens: 1,
54. 54 output_tokens: 1,
55. 55 ..Default::default()
56. 56 },
57. 57 model: "test-model".into(),
58. 58 response_id: None,
59. 59 }
60. 60}
61. 61
62. 62fn bash_call(id: &str, cmd: &str) -> AssistantMessage {
63. 63 AssistantMessage {
64. 64 content: vec![ContentBlock::ToolUse {
65. 65 id: id.into(),
66. 66 name: "bash".into(),
67. 67 input: serde_json::json!({"command": cmd}),
68. 68 }],
69. 69 stop_reason: StopReason::ToolUse,
70. 70 usage: Usage::default(),
71. 71 model: "test-model".into(),
72. 72 response_id: None,
73. 73 }
74. 74}
75. 75
76. 76fn build(responses: Vec<AssistantMessage>, hooks: Option<Vec<HookDef>>) -> Agent {
77. 77 let dir = tempdir().unwrap();
78. 78 let header = SessionHeader {
79. 79 agent: None,
80. 80 session_id: "s-hooks".into(),
81. 81 created_at: chrono::Utc::now(),
82. 82 cwd: dir.path().to_path_buf(),
83. 83 parent_session_id: None,
84. 84 contract_id: None,
85. 85 work_item_id: None,
86. 86 conversation: None,
87. 87 contract: FrozenContract {
88. 88 app_version: "0".into(),
89. 89 provider: "scripted".into(),
90. 90 model: "test-model".into(),
91. 91 route_ladder: Vec::new(),
92. 92 route_objective: String::new(),
93. 93 route_annotations: Vec::new(),
94. 94 system_prompt: "sys".into(),
95. 95 permission_mode: "full-access".into(),
96. 96 capabilities: Vec::new(),
97. 97 prompt_layers: Vec::new(),
98. 98 },
99. 99 };
100. 100 let log = SessionLog::create(dir.path().join("s.jsonl"), header).unwrap();
101. 101 let mut cfg = AgentConfig::new("sys");
102. 102 cfg.tools = vec![Arc::new(BashTool)];
103. 103 cfg.mode = vak_permission::Mode::FullAccess;
104. 104 cfg.hooks = hooks.map(Arc::new);
105. 105 cfg.stop_policy = None;
106. 106 std::mem::forget(dir);
107. 107 Agent::new(
108. 108 Arc::new(Scripted {
109. 109 responses: Mutex::new(responses.into_iter().collect()),
110. 110 }),
111. 111 log,
112. 112 cfg,
113. 113 )
114. 114}
115. 115
116. 116#[tokio::test]
117. 117async fn pre_tool_use_hook_blocks_execution() {
118. 118 let marker = tempdir().unwrap();
119. 119 let marker_path = marker.path().join("ran");
120. 120 let hooks = vec![HookDef {
121. 121 event: vak_hooks::HookEvent::PreToolUse,
122. 122 matcher: Some(vak_permission::Rule::parse("Bash(touch *)").unwrap()),
123. 123 command: r#"echo '{"decision":"block","reason":"no touching"}'"#.to_string(),
124. 124 timeout_ms: 5000,
125. 125 failure_mode: vak_hooks::HookFailureMode::Open,
126. 126 refusal: None,
127. 127 }];
128. 128 let marker_path = marker_path.display().to_string();
129. 129 let mut agent = build(
130. 130 vec![
131. 131 bash_call("t1", &format!("touch {marker_path}")),
132. 132 text_msg("adapted"),
133. 133 ],
134. 134 Some(hooks),
135. 135 );
136. 136
137. 137 let outcome = agent
138. 138 .run(
139. 139 "go",
140. 140 &Default::default(),
141. 141 CancellationToken::new(),
142. 142 mpsc::channel(64).0,
143. 143 )
144. 144 .await;
145. 145 assert!(matches!(outcome, TurnOutcome::Completed { .. }));
146. 146 assert!(
147. 147 !std::path::Path::new(&marker_path).exists(),
148. 148 "blocked command must never have executed"
149. 149 );
150. 150 let session = agent.session.lock().await;
151. 151 // Raw ledger: the closed turn's result is a trace line in the
152. 152 // projection now (docs/design/68-context-engine.md §10).
153. 153 let result = session
154. 154 .message_chain()
155. 155 .iter()
156. 156 .flat_map(|(_, m)| m.content.iter())
157. 157 .find_map(|b| match b {
158. 158 ContentBlock::ToolResult {
159. 159 content, is_error, ..
160. 160 } => Some((content.clone(), *is_error)),
161. 161 _ => None,
162. 162 })
163. 163 .unwrap();
164. 164 assert!(result.1);
165. 165 assert!(result.0.contains("no touching"));
166. 166}
167. 167
168. 168#[tokio::test]
169. 169async fn stop_hook_forces_continuation_once() {
170. 170 let dir = tempdir().unwrap();
171. 171 let flag = dir.path().join("block-once");
172. 172 let flag_display = flag.display().to_string();
173. 173 // Block the first Stop; after the flag exists, allow it.
174. 174 let cmd = format!(
175. 175 "if [ -f {flag_display} ]; then exit 0; else touch {flag_display}; echo '{{\"decision\":\"block\",\"reason\":\"say goodbye first\"}}'; fi"
176. 176 );
177. 177 let hooks = vec![HookDef {
178. 178 event: vak_hooks::HookEvent::Stop,
179. 179 matcher: None,
180. 180 command: cmd,
181. 181 timeout_ms: 5000,
182. 182 failure_mode: vak_hooks::HookFailureMode::Open,
183. 183 refusal: None,
184. 184 }];
185. 185 let mut agent = build(
186. 186 vec![text_msg("first attempt"), text_msg("goodbye")],
187. 187 Some(hooks),
188. 188 );
189. 189
190. 190 let outcome = agent
191. 191 .run(
192. 192 "go",
193. 193 &Default::default(),
194. 194 CancellationToken::new(),
195. 195 mpsc::channel(64).0,
196. 196 )
197. 197 .await;
198. 198 match outcome {
199. 199 TurnOutcome::Completed { response } => {
200. 200 assert_eq!(response.text_content(), "goodbye");
201. 201 }
202. 202 other => panic!("expected completed after continuation, got {other:?}"),
203. 203 }
204. 204 let session = agent.session.lock().await;
205. 205 // Raw ledger: a control nudge is scaffolding for the turn still in
206. 206 // progress and is dropped once the turn closes
207. 207 // (docs/design/68-context-engine.md §10).
208. 208 let texts: Vec<String> = session
209. 209 .message_chain()
210. 210 .iter()
211. 211 .filter(|(_, m)| m.text_content().contains("[stop-hook]"))
212. 212 .map(|(_, m)| m.text_content())
213. 213 .collect();
214. 214 assert_eq!(texts.len(), 1, "continuation message must be logged once");
215. 215}
216. 216