1. 1//! End-to-end guards for the capability lifecycle
2. 2//! (docs/design/41-capability-registry.md).
3. 3//!
4. 4//! Every test here pins a behaviour that was broken in a way no error
5. 5//! surfaced: the agent did not fail, it answered from memory and sounded
6. 6//! certain. That invisibility is why these are integration tests rather than
7. 7//! notes in a changelog.
8. 8
9. 9#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
10. 10
11. 11use std::collections::{BTreeSet, VecDeque};
12. 12use std::path::Path;
13. 13use std::sync::{Arc, Mutex};
14. 14
15. 15use tokio_util::sync::CancellationToken;
16. 16
17. 17use vak_core::Core;
18. 18use vak_core::capability::registry::{CapabilityProvider, Declaration};
19. 19use vak_core::capability::{CapabilityId, CapabilityRegistry, Domain, Origin, Serves};
20. 20use vak_llm::stream;
21. 21use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, Usage};
22. 22use vak_llm::{EventStream, LlmError, Provider};
23. 23use vak_session::types::CapabilityKind;
24. 24
25. 25// ---------------------------------------------------------------- fakes ---
26. 26
27. 27struct Fake {
28. 28 declarations: Mutex<Vec<Declaration>>,
29. 29}
30. 30
31. 31impl Fake {
32. 32 fn new(declarations: Vec<Declaration>) -> Arc<Self> {
33. 33 Arc::new(Fake {
34. 34 declarations: Mutex::new(declarations),
35. 35 })
36. 36 }
37. 37}
38. 38
39. 39#[async_trait::async_trait]
40. 40impl CapabilityProvider for Fake {
41. 41 fn declare(&self) -> Vec<Declaration> {
42. 42 self.declarations.lock().unwrap().clone()
43. 43 }
44. 44}
45. 45
46. 46fn declaration(name: &str, kind: CapabilityKind) -> Declaration {
47. 47 Declaration {
48. 48 id: CapabilityId::new(kind, name),
49. 49 origin: Origin::Workspace,
50. 50 summary: String::new(),
51. 51 serves: Serves::Undeclared,
52. 52 digest: None,
53. 53 source: None,
54. 54 configuration: serde_json::Value::Null,
55. 55 }
56. 56}
57. 57
58. 58struct Scripted {
59. 59 responses: Mutex<VecDeque<AssistantMessage>>,
60. 60}
61. 61
62. 62#[async_trait::async_trait]
63. 63impl Provider for Scripted {
64. 64 fn name(&self) -> &str {
65. 65 "scripted"
66. 66 }
67. 67
68. 68 async fn stream(
69. 69 &self,
70. 70 _request: ChatRequest,
71. 71 _cancel: CancellationToken,
72. 72 ) -> Result<EventStream, LlmError> {
73. 73 let next = self.responses.lock().unwrap().pop_front();
74. 74 let (mut sink, rx) = stream::channel(64);
75. 75 match next {
76. 76 Some(m) => {
77. 77 sink.push(stream::StreamEvent::Start { partial: m.clone() });
78. 78 sink.close_message(m).await;
79. 79 }
80. 80 None => sink.close_error(LlmError::Parse("exhausted".into())).await,
81. 81 }
82. 82 Ok(rx)
83. 83 }
84. 84}
85. 85
86. 86fn reply(content: ContentBlock, stop: vak_llm::types::StopReason) -> AssistantMessage {
87. 87 AssistantMessage {
88. 88 content: vec![content],
89. 89 stop_reason: stop,
90. 90 usage: Usage::default(),
91. 91 model: "test-model".into(),
92. 92 response_id: None,
93. 93 }
94. 94}
95. 95
96. 96fn spawns(marker: &Path) -> usize {
97. 97 std::fs::read_to_string(marker)
98. 98 .map(|text| text.lines().count())
99. 99 .unwrap_or(0)
100. 100}
101. 101
102. 102// ------------------------------------------------------------- the case ---
103. 103
104. 104/// A configured MCP server is started by demand and nothing else
105. 105/// (AGENTS.md invariant 25). Admission, reconciliation and prompt assembly
106. 106/// all run without starting it; the model's first `mcp` call starts it once;
107. 107/// and what that call learned reaches the next turn's prompt with no restart.
108. 108#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
109. 109async fn an_mcp_server_starts_on_demand_and_its_catalog_reaches_the_next_turn() {
110. 110 vak_config::paths::isolate_home_for_tests();
111. 111 let dir = tempfile::tempdir().unwrap();
112. 112 let cwd = dir.path().join("workspace");
113. 113 std::fs::create_dir_all(&cwd).unwrap();
114. 114 let marker = dir.path().join("spawns");
115. 115 let script = Path::new(env!("CARGO_MANIFEST_DIR"))
116. 116 .join("../../scripts/fake_mcp_server.py")
117. 117 .canonicalize()
118. 118 .unwrap();
119. 119
120. 120 let core = Core::new_with_trust(cwd, true).unwrap();
121. 121 core.set_sessions_home(dir.path().join("home"));
122. 122 core.set_permission_mode(vak_config::PermissionMode::FullAccess);
123. 123 let mut mcp = vak_config::McpConfig::default();
124. 124 mcp.servers.insert(
125. 125 "fake".into(),
126. 126 vak_config::McpServerConfig {
127. 127 command: "sh".into(),
128. 128 args: vec![
129. 129 "-c".into(),
130. 130 format!(
131. 131 "echo started >> '{}'; exec python3 '{}'",
132. 132 marker.display(),
133. 133 script.display()
134. 134 ),
135. 135 ],
136. 136 env: Default::default(),
137. 137 network: false,
138. 138 serves: Vec::new(),
139. 139 },
140. 140 );
141. 141 core.set_mcp_servers(mcp);
142. 142
143. 143 core.admitted_capabilities().await;
144. 144 core.reconcile_capabilities().await;
145. 145 let before = core.system_prompt();
146. 146 assert!(before.contains("\n- fake\n"), "named, not yet listed");
147. 147 assert_eq!(
148. 148 spawns(&marker),
149. 149 0,
150. 150 "admission and prompt assembly start nothing"
151. 151 );
152. 152
153. 153 core.set_provider_instance(Arc::new(Scripted {
154. 154 responses: Mutex::new(VecDeque::from(vec![
155. 155 reply(
156. 156 ContentBlock::ToolUse {
157. 157 id: "m1".into(),
158. 158 name: "mcp".into(),
159. 159 input: serde_json::json!({"action": "list", "server": "fake"}),
160. 160 },
161. 161 vak_llm::types::StopReason::ToolUse,
162. 162 ),
163. 163 reply(
164. 164 ContentBlock::text("listed"),
165. 165 vak_llm::types::StopReason::EndTurn,
166. 166 ),
167. 167 ])),
168. 168 }));
169. 169 let session = core.start_session().await.unwrap();
170. 170 let (events, _rx) = tokio::sync::mpsc::channel(256);
171. 171 let (outcome, _) = core
172. 172 .run_turn_with(
173. 173 session,
174. 174 "what can the fake server do?",
175. 175 CancellationToken::new(),
176. 176 None,
177. 177 None,
178. 178 None,
179. 179 events,
180. 180 )
181. 181 .await
182. 182 .unwrap();
183. 183 assert!(matches!(outcome, vak_agent::TurnOutcome::Completed { .. }));
184. 184 assert_eq!(spawns(&marker), 1, "the model's call started it, once");
185. 185
186. 186 core.reconcile_capabilities().await;
187. 187 assert!(
188. 188 core.system_prompt()
189. 189 .contains("- fake: echo, boom, secret_result"),
190. 190 "the observed catalog reaches the next prompt: {}",
191. 191 core.system_prompt()
192. 192 );
193. 193 assert_eq!(
194. 194 spawns(&marker),
195. 195 1,
196. 196 "learning the catalog started nothing new"
197. 197 );
198. 198}
199. 199
200. 200// ------------------------------------------------------------ lifecycle ---
201. 201
202. 202/// Add, edit and remove, on a registry that is never restarted.
203. 203#[tokio::test]
204. 204async fn capabilities_can_be_added_edited_and_removed_while_running() {
205. 205 let provider = Fake::new(vec![declaration("read", CapabilityKind::Tool)]);
206. 206 let (registry, _hints) = CapabilityRegistry::new(provider.clone());
207. 207 registry.reconcile().await;
208. 208 let first = registry.current().await.epoch;
209. 209
210. 210 // --- added -----------------------------------------------------------
211. 211 provider
212. 212 .declarations
213. 213 .lock()
214. 214 .unwrap()
215. 215 .push(declaration("pdf", CapabilityKind::Skill));
216. 216 let delta = registry.reconcile().await.expect("adding is a change");
217. 217 assert_eq!(delta.added.len(), 1);
218. 218 assert!(delta.describe().contains("now available"));
219. 219 let second = registry.current().await.epoch;
220. 220 assert!(second > first, "a change publishes a new epoch");
221. 221
222. 222 // --- edited ----------------------------------------------------------
223. 223 provider
224. 224 .declarations
225. 225 .lock()
226. 226 .unwrap()
227. 227 .iter_mut()
228. 228 .filter(|d| d.id.name == "pdf")
229. 229 .for_each(|d| d.digest = Some("rewritten".into()));
230. 230 let delta = registry.reconcile().await.expect("editing is a change");
231. 231 assert_eq!(delta.updated.len(), 1);
232. 232 assert!(registry.current().await.epoch > second);
233. 233
234. 234 // --- removed ---------------------------------------------------------
235. 235 provider
236. 236 .declarations
237. 237 .lock()
238. 238 .unwrap()
239. 239 .retain(|d| d.id.name != "pdf");
240. 240 let delta = registry.reconcile().await.expect("removing is a change");
241. 241 assert_eq!(delta.removed.len(), 1);
242. 242 assert!(
243. 243 delta.describe().contains("no longer available"),
244. 244 "removal is announced, not silent"
245. 245 );
246. 246}
247. 247
248. 248/// Reconciling an unchanged world must publish nothing.
249. 249///
250. 250/// A loop that ran every ten seconds for three weeks and republished each
251. 251/// time would churn every live session's prompt and defeat the point.
252. 252#[tokio::test]
253. 253async fn a_quiet_world_never_churns_the_epoch() {
254. 254 let provider = Fake::new(vec![declaration("read", CapabilityKind::Tool)]);
255. 255 let (registry, _hints) = CapabilityRegistry::new(provider);
256. 256 registry.reconcile().await;
257. 257 let epoch = registry.current().await.epoch;
258. 258 for _ in 0..25 {
259. 259 assert!(registry.reconcile().await.is_none());
260. 260 }
261. 261 assert_eq!(registry.current().await.epoch, epoch);
262. 262}
263. 263
264. 264// ----------------------------------------------------------- revocation ---
265. 265
266. 266/// Revocation is immediate and independent of any epoch. An operator
267. 267/// disabling a compromised plugin must not wait for a long turn to finish.
268. 268#[tokio::test]
269. 269async fn revocation_applies_immediately_and_restoration_republishes() {
270. 270 let provider = Fake::new(vec![declaration("bash", CapabilityKind::Tool)]);
271. 271 let (registry, _hints) = CapabilityRegistry::new(provider);
272. 272 registry.reconcile().await;
273. 273 let id = CapabilityId::new(CapabilityKind::Tool, "bash");
274. 274
275. 275 registry.revoke(id.clone(), "operator disabled").await;
276. 276 assert_eq!(
277. 277 registry.revocation(&id).await.as_deref(),
278. 278 Some("operator disabled"),
279. 279 "visible at dispatch with no reconcile in between"
280. 280 );
281. 281
282. 282 registry.reconcile().await;
283. 283 assert_eq!(registry.current().await.usable().count(), 0);
284. 284
285. 285 registry.restore(&id).await;
286. 286 registry.reconcile().await;
287. 287 assert_eq!(registry.current().await.usable().count(), 1);
288. 288}
289. 289
290. 290// -------------------------------------------------------------- domains ---
291. 291
292. 292/// A domain this build has never heard of must round-trip and match itself,
293. 293/// so a plugin with its own vocabulary is not silently flattened.
294. 294#[test]
295. 295fn an_unknown_domain_survives_and_matches() {
296. 296 let parsed = Domain::parse("procurement");
297. 297 assert_eq!(parsed, Domain::Custom("procurement".into()));
298. 298 let serves = Serves::declared([parsed.clone()]);
299. 299 assert!(serves.serves_any(&BTreeSet::from([parsed])));
300. 300 assert!(!serves.serves_any(&BTreeSet::from([Domain::Web])));
301. 301}
302. 302
303. 303/// A changed secret changes the server's declared digest, which is a change
304. 304/// to the world: admission sees pending work and the next turn binds the new
305. 305/// epoch, with no restart.
306. 306#[tokio::test]
307. 307async fn a_resolved_secret_is_a_pending_change_at_the_next_admission() {
308. 308 let mut decl = declaration("search_srv", CapabilityKind::McpServer);
309. 309 decl.digest = Some("unresolved_hash".into());
310. 310 let provider = Fake::new(vec![decl.clone()]);
311. 311 let (registry, _hints) = CapabilityRegistry::new(provider.clone());
312. 312 registry.reconcile().await;
313. 313 let first = registry.current().await.epoch;
314. 314
315. 315 decl.digest = Some("resolved_hash_with_key".into());
316. 316 decl.configuration = serde_json::json!({"tools": [{"name": "search"}]});
317. 317 *provider.declarations.lock().unwrap() = vec![decl];
318. 318 assert!(registry.has_pending_changes().await);
319. 319 assert!(registry.reconcile().await.is_some());
320. 320 let current = registry.current().await;
321. 321 assert!(current.epoch > first);
322. 322
323. 323 let inventory = current.mcp_inventory();
324. 324 let empty_policy = vak_config::ChannelPolicy::default();
325. 325 let builtins: BTreeSet<String> = ["read".to_string(), "glob".to_string()].into();
326. 326 let tc = vak_core::capability::TurnCapabilities::build(&vak_core::capability::TurnProbe {
327. 327 capabilities: &current,
328. 328 revoked_ids: BTreeSet::new(),
329. 329 channel_policy: &empty_policy,
330. 330 reach_standings: &[],
331. 331 mcp_inventory: &inventory,
332. 332 builtin_names: &builtins,
333. 333 });
334. 334 assert_eq!(
335. 335 tc.mcp_tool_index.get("search").map(String::as_str),
336. 336 Some("search_srv"),
337. 337 "a bare `search` call resolves to its server in turn 2"
338. 338 );
339. 339 assert!(tc.mcp_server_names.contains(&"search_srv".to_string()));
340. 340}
341. 341
342. 342// -------------------------------------------------------- agent network ---
343. 343
344. 344/// Inter-agent messaging is offered only while an operator has authorized
345. 345/// this workspace on the broker, and withdrawn when that authorization is.
346. 346#[tokio::test]
347. 347async fn agent_network_is_offered_only_while_the_workspace_is_authorized() {
348. 348 vak_config::paths::isolate_home_for_tests();
349. 349 let dir = tempfile::tempdir().unwrap();
350. 350 let core = Core::new_with_trust(dir.path().to_path_buf(), true).unwrap();
351. 351 let offered = |core: &Core| core.tool_names().iter().any(|n| n == "agent_network");
352. 352 assert!(!offered(&core), "off by default");
353. 353
354. 354 let broker = core.agent_network_broker();
355. 355 let capability = broker.register(
356. 356 core.cwd().canonicalize().unwrap().display().to_string(),
357. 357 vak_core::agent_network::WorkspaceNetworkPolicy {
358. 358 enabled: true,
359. 359 allowed_peers: ["peer-workspace".to_string()].into(),
360. 360 max_message_bytes: 1024,
361. 361 },
362. 362 );
363. 363 assert!(offered(&core));
364. 364
365. 365 assert!(broker.revoke(&capability));
366. 366 assert!(
367. 367 !offered(&core),
368. 368 "revoking the registration withdraws the tool"
369. 369 );
370. 370}
371. 371