1. 1#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
2. 2
3. 3//! Doom-loop guard: the third identical (tool, args) call in one run is
4. 4//! re-routed through approval instead of executing silently.
5. 5
6. 6use std::collections::VecDeque;
7. 7use std::sync::{Arc, Mutex};
8. 8
9. 9use tokio::sync::mpsc;
10. 10use tokio_util::sync::CancellationToken;
11. 11
12. 12use tempfile::tempdir;
13. 13
14. 14use vak_agent::{Agent, AgentConfig, TurnOutcome};
15. 15use vak_llm::stream;
16. 16use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, StopReason, Usage};
17. 17use vak_llm::{EventStream, LlmError, Provider};
18. 18use vak_permission::PermissionEngine;
19. 19use vak_session::types::{FrozenContract, SessionHeader};
20. 20use vak_session::{SessionLog, SessionPath};
21. 21use vak_tools::bash::BashTool;
22. 22
23. 23struct Scripted {
24. 24 responses: Mutex<VecDeque<AssistantMessage>>,
25. 25}
26. 26
27. 27#[async_trait::async_trait]
28. 28impl Provider for Scripted {
29. 29 fn name(&self) -> &str {
30. 30 "scripted"
31. 31 }
32. 32
33. 33 async fn stream(
34. 34 &self,
35. 35 _request: ChatRequest,
36. 36 _cancel: CancellationToken,
37. 37 ) -> Result<EventStream, LlmError> {
38. 38 let next = self.responses.lock().unwrap().pop_front();
39. 39 let (mut sink, rx) = stream::channel(64);
40. 40 match next {
41. 41 Some(m) => {
42. 42 sink.push(stream::StreamEvent::Start { partial: m.clone() });
43. 43 sink.close_message(m).await;
44. 44 }
45. 45 None => sink.close_error(LlmError::Parse("exhausted".into())).await,
46. 46 }
47. 47 Ok(rx)
48. 48 }
49. 49}
50. 50
51. 51fn bash_call(id: &str, cmd: &str) -> AssistantMessage {
52. 52 AssistantMessage {
53. 53 content: vec![ContentBlock::ToolUse {
54. 54 id: id.into(),
55. 55 name: "bash".into(),
56. 56 input: serde_json::json!({"command": cmd}),
57. 57 }],
58. 58 stop_reason: StopReason::ToolUse,
59. 59 usage: Usage::default(),
60. 60 model: "test-model".into(),
61. 61 response_id: None,
62. 62 }
63. 63}
64. 64
65. 65fn text_msg(t: &str) -> AssistantMessage {
66. 66 AssistantMessage {
67. 67 content: vec![ContentBlock::text(t)],
68. 68 stop_reason: StopReason::EndTurn,
69. 69 usage: Usage {
70. 70 input_tokens: 1,
71. 71 output_tokens: 1,
72. 72 ..Default::default()
73. 73 },
74. 74 model: "test-model".into(),
75. 75 response_id: None,
76. 76 }
77. 77}
78. 78
79. 79#[tokio::test]
80. 80async fn third_identical_call_is_blocked_with_reason() {
81. 81 let dir = tempdir().unwrap();
82. 82 let header = SessionHeader {
83. 83 agent: None,
84. 84 session_id: "doom-loop".into(),
85. 85 created_at: chrono::Utc::now(),
86. 86 cwd: dir.path().to_path_buf(),
87. 87 parent_session_id: None,
88. 88 contract_id: None,
89. 89 work_item_id: None,
90. 90 conversation: None,
91. 91 contract: FrozenContract {
92. 92 app_version: "0".into(),
93. 93 provider: "scripted".into(),
94. 94 model: "test-model".into(),
95. 95 route_ladder: Vec::new(),
96. 96 route_objective: String::new(),
97. 97 route_annotations: Vec::new(),
98. 98 system_prompt: "sys".into(),
99. 99 permission_mode: "full-access".into(),
100. 100 capabilities: Vec::new(),
101. 101 prompt_layers: Vec::new(),
102. 102 },
103. 103 };
104. 104 let home = dir.path().join("home");
105. 105 std::fs::create_dir_all(&home).unwrap();
106. 106 let log = SessionLog::create(
107. 107 SessionPath::new_session_file(&home, dir.path(), "doom-loop"),
108. 108 header,
109. 109 )
110. 110 .unwrap();
111. 111
112. 112 let mut agent = Agent::new(
113. 113 Arc::new(Scripted {
114. 114 responses: Mutex::new(VecDeque::from(vec![
115. 115 bash_call("a", "echo same"),
116. 116 bash_call("b", "echo same"),
117. 117 bash_call("c", "echo same"),
118. 118 text_msg("adapted after guard"),
119. 119 ])),
120. 120 }),
121. 121 log,
122. 122 {
123. 123 let mut cfg = AgentConfig::new("sys");
124. 124 cfg.model = "test-model".into();
125. 125 cfg.tools = vec![Arc::new(BashTool)];
126. 126 cfg.mode = vak_permission::Mode::FullAccess;
127. 127 cfg.permission = Some(Arc::new(PermissionEngine::default()));
128. 128 // AutoDeny makes the doom-loop Ask observable as a typed denial.
129. 129 cfg.approver = Some(Arc::new(vak_agent::AutoDeny));
130. 130 cfg
131. 131 },
132. 132 );
133. 133
134. 134 let outcome = agent
135. 135 .run(
136. 136 "loop",
137. 137 &Default::default(),
138. 138 CancellationToken::new(),
139. 139 mpsc::channel(64).0,
140. 140 )
141. 141 .await;
142. 142 assert!(
143. 143 matches!(outcome, TurnOutcome::Completed { .. }),
144. 144 "got {outcome:?}"
145. 145 );
146. 146
147. 147 // Raw ledger: the closed turn's results are trace lines in the
148. 148 // model-visible projection now (docs/design/68-context-engine.md
149. 149 // §10); this test is about what the guard actually recorded.
150. 150 let results: Vec<(bool, String)> = agent
151. 151 .session
152. 152 .lock()
153. 153 .await
154. 154 .message_chain()
155. 155 .iter()
156. 156 .flat_map(|(_, m)| m.content.iter())
157. 157 .filter_map(|b| match b {
158. 158 ContentBlock::ToolResult {
159. 159 content, is_error, ..
160. 160 } => Some((*is_error, content.clone())),
161. 161 _ => None,
162. 162 })
163. 163 .collect();
164. 164 assert_eq!(results.len(), 3);
165. 165 assert!(!results[0].0 && !results[1].0, "first two run normally");
166. 166 assert!(results[2].0, "third identical call must be blocked");
167. 167 assert!(
168. 168 results[2].1.contains("repeated"),
169. 169 "guard reason must reach the model: {}",
170. 170 results[2].1
171. 171 );
172. 172}
173. 173
174. 174#[tokio::test]
175. 175async fn different_args_are_not_counted_together() {
176. 176 let dir = tempdir().unwrap();
177. 177 let header = SessionHeader {
178. 178 agent: None,
179. 179 session_id: "no-doom".into(),
180. 180 created_at: chrono::Utc::now(),
181. 181 cwd: dir.path().to_path_buf(),
182. 182 parent_session_id: None,
183. 183 contract_id: None,
184. 184 work_item_id: None,
185. 185 conversation: None,
186. 186 contract: FrozenContract {
187. 187 app_version: "0".into(),
188. 188 provider: "scripted".into(),
189. 189 model: "test-model".into(),
190. 190 route_ladder: Vec::new(),
191. 191 route_objective: String::new(),
192. 192 route_annotations: Vec::new(),
193. 193 system_prompt: "sys".into(),
194. 194 permission_mode: "full-access".into(),
195. 195 capabilities: Vec::new(),
196. 196 prompt_layers: Vec::new(),
197. 197 },
198. 198 };
199. 199 let home = dir.path().join("home");
200. 200 std::fs::create_dir_all(&home).unwrap();
201. 201 let log = SessionLog::create(
202. 202 SessionPath::new_session_file(&home, dir.path(), "no-doom"),
203. 203 header,
204. 204 )
205. 205 .unwrap();
206. 206
207. 207 let mut agent = Agent::new(
208. 208 Arc::new(Scripted {
209. 209 responses: Mutex::new(VecDeque::from(vec![
210. 210 bash_call("a", "echo one"),
211. 211 bash_call("b", "echo two"),
212. 212 bash_call("c", "echo three"),
213. 213 text_msg("done"),
214. 214 ])),
215. 215 }),
216. 216 log,
217. 217 {
218. 218 let mut cfg = AgentConfig::new("sys");
219. 219 cfg.model = "test-model".into();
220. 220 cfg.tools = vec![Arc::new(BashTool)];
221. 221 cfg.mode = vak_permission::Mode::FullAccess;
222. 222 cfg.permission = Some(Arc::new(PermissionEngine::default()));
223. 223 cfg.approver = Some(Arc::new(vak_agent::AutoDeny));
224. 224 cfg
225. 225 },
226. 226 );
227. 227
228. 228 let outcome = agent
229. 229 .run(
230. 230 "distinct commands",
231. 231 &Default::default(),
232. 232 CancellationToken::new(),
233. 233 mpsc::channel(64).0,
234. 234 )
235. 235 .await;
236. 236 assert!(matches!(outcome, TurnOutcome::Completed { .. }));
237. 237 for is_error in agent
238. 238 .session
239. 239 .lock()
240. 240 .await
241. 241 .derive_messages()
242. 242 .iter()
243. 243 .flat_map(|m| m.content.iter())
244. 244 .filter_map(|b| match b {
245. 245 ContentBlock::ToolResult { is_error, .. } => Some(*is_error),
246. 246 _ => None,
247. 247 })
248. 248 {
249. 249 assert!(!is_error, "distinct calls must never trip the guard");
250. 250 }
251. 251}
252. 252