1. 1#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
2. 2
3. 3//! Regression coverage for the verified real bug: a retrieval-shaped tool
4. 4//! call (e.g. `tavily_search`) succeeds and returns real results, but the
5. 5//! model's next turn ignores them and answers with ungrounded, uncited
6. 6//! prose. `Agent::run` must give the model exactly one bounded repair turn
7. 7//! in that case instead of letting the ungrounded answer stand.
8. 8//!
9. 9//! Which calls count as retrieval is decided by `AgentConfig::retrieval_check`,
10. 10//! supplied by `Core` from what each capability *declares it serves* — never
11. 11//! from a tool's name or its output. These tests supply a check that declares
12. 12//! the tool `search` as web-serving, and separately prove the name is
13. 13//! irrelevant in both directions.
14. 14
15. 15use std::collections::VecDeque;
16. 16use std::sync::{Arc, Mutex};
17. 17
18. 18use async_trait::async_trait;
19. 19use serde_json::Value;
20. 20use tokio::sync::mpsc;
21. 21use tokio_util::sync::CancellationToken;
22. 22
23. 23use tempfile::tempdir;
24. 24
25. 25use vak_agent::{Agent, AgentConfig, TurnOutcome};
26. 26use vak_llm::stream;
27. 27use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, StopReason, Usage};
28. 28use vak_llm::{EventStream, LlmError, Provider};
29. 29use vak_permission::PermissionEngine;
30. 30use vak_session::types::{FrozenContract, SessionHeader};
31. 31use vak_session::{SessionLog, SessionPath};
32. 32use vak_tools::context::ToolContext;
33. 33use vak_tools::{Tool, ToolOutput};
34. 34
35. 35struct Scripted {
36. 36 responses: Mutex<VecDeque<AssistantMessage>>,
37. 37}
38. 38
39. 39#[async_trait]
40. 40impl Provider for Scripted {
41. 41 fn name(&self) -> &str {
42. 42 "scripted"
43. 43 }
44. 44
45. 45 async fn stream(
46. 46 &self,
47. 47 _request: ChatRequest,
48. 48 _cancel: CancellationToken,
49. 49 ) -> Result<EventStream, LlmError> {
50. 50 let next = self.responses.lock().unwrap().pop_front();
51. 51 let (mut sink, rx) = stream::channel(64);
52. 52 match next {
53. 53 Some(m) => {
54. 54 sink.push(stream::StreamEvent::Start { partial: m.clone() });
55. 55 sink.close_message(m).await;
56. 56 }
57. 57 None => sink.close_error(LlmError::Parse("exhausted".into())).await,
58. 58 }
59. 59 Ok(rx)
60. 60 }
61. 61}
62. 62
63. 63/// Stands in for any retrieval-type MCP tool (Tavily/Exa/Firecrawl/a future
64. 64/// one nobody's written yet) — real, dated, URL-bearing results, exactly
65. 65/// the shape observed in the live bug's session transcript.
66. 66struct FakeSearchTool;
67. 67
68. 68#[async_trait]
69. 69impl Tool for FakeSearchTool {
70. 70 fn name(&self) -> &str {
71. 71 "search"
72. 72 }
73. 73 fn description(&self) -> &str {
74. 74 "fake search tool for tests"
75. 75 }
76. 76 fn schema(&self) -> Value {
77. 77 serde_json::json!({"type": "object", "properties": {"query": {"type": "string"}}})
78. 78 }
79. 79 async fn execute(&self, _args: &Value, _ctx: &ToolContext) -> ToolOutput {
80. 80 ToolOutput::ok(
81. 81 "Title: Parliament debates UPI fee proposal\nURL: https://example-news.in/upi-fee-debate\n\n\
82. 82 Title: RBI holds rates steady\nURL: https://example-news.in/rbi-rates\n",
83. 83 )
84. 84 }
85. 85}
86. 86
87. 87fn search_call(id: &str) -> AssistantMessage {
88. 88 AssistantMessage {
89. 89 content: vec![ContentBlock::ToolUse {
90. 90 id: id.into(),
91. 91 name: "search".into(),
92. 92 input: serde_json::json!({"query": "top news in india right now"}),
93. 93 }],
94. 94 stop_reason: StopReason::ToolUse,
95. 95 usage: Usage::default(),
96. 96 model: "test-model".into(),
97. 97 response_id: None,
98. 98 }
99. 99}
100. 100
101. 101fn text_msg(t: &str) -> AssistantMessage {
102. 102 AssistantMessage {
103. 103 content: vec![ContentBlock::text(t)],
104. 104 stop_reason: StopReason::EndTurn,
105. 105 usage: Usage {
106. 106 input_tokens: 1,
107. 107 output_tokens: 1,
108. 108 ..Default::default()
109. 109 },
110. 110 model: "test-model".into(),
111. 111 response_id: None,
112. 112 }
113. 113}
114. 114
115. 115fn fenced_research_answer() -> String {
116. 116 "Here is what I found:\n\n```vak\n{\"semantic_type\":\"research.synthesis\",\"payload\":{\"sources\":[{\"title\":\"UPI fee debate\",\"url\":\"https://example-news.in/upi-fee-debate\"}],\"takeaways\":[{\"text\":\"Parliament is debating a UPI fee.\",\"citation_indices\":[0]}]}}\n```".into()
117. 117}
118. 118
119. 119async fn build_agent(
120. 120 dir: &tempfile::TempDir,
121. 121 session_id: &str,
122. 122 responses: Vec<AssistantMessage>,
123. 123) -> Agent {
124. 124 build_agent_with_check(
125. 125 dir,
126. 126 session_id,
127. 127 responses,
128. 128 Some(Arc::new(|name: &str, _: &Value| name == "search")),
129. 129 )
130. 130 .await
131. 131}
132. 132
133. 133async fn build_agent_with_check(
134. 134 dir: &tempfile::TempDir,
135. 135 session_id: &str,
136. 136 responses: Vec<AssistantMessage>,
137. 137 retrieval_check: Option<vak_agent::RetrievalCheck>,
138. 138) -> Agent {
139. 139 let header = SessionHeader {
140. 140 agent: None,
141. 141 session_id: session_id.into(),
142. 142 created_at: chrono::Utc::now(),
143. 143 cwd: dir.path().to_path_buf(),
144. 144 parent_session_id: None,
145. 145 contract_id: None,
146. 146 work_item_id: None,
147. 147 conversation: None,
148. 148 contract: FrozenContract {
149. 149 app_version: "0".into(),
150. 150 provider: "scripted".into(),
151. 151 model: "test-model".into(),
152. 152 route_ladder: Vec::new(),
153. 153 route_objective: String::new(),
154. 154 route_annotations: Vec::new(),
155. 155 system_prompt: "sys".into(),
156. 156 permission_mode: "full-access".into(),
157. 157 capabilities: Vec::new(),
158. 158 prompt_layers: Vec::new(),
159. 159 },
160. 160 };
161. 161 let home = dir.path().join("home");
162. 162 std::fs::create_dir_all(&home).unwrap();
163. 163 let log = SessionLog::create(
164. 164 SessionPath::new_session_file(&home, dir.path(), session_id),
165. 165 header,
166. 166 )
167. 167 .unwrap();
168. 168
169. 169 Agent::new(
170. 170 Arc::new(Scripted {
171. 171 responses: Mutex::new(VecDeque::from(responses)),
172. 172 }),
173. 173 log,
174. 174 {
175. 175 let mut cfg = AgentConfig::new("sys");
176. 176 cfg.model = "test-model".into();
177. 177 cfg.tools = vec![Arc::new(FakeSearchTool)];
178. 178 cfg.retrieval_check = retrieval_check;
179. 179 cfg.mode = vak_permission::Mode::FullAccess;
180. 180 cfg.permission = Some(Arc::new(PermissionEngine::default()));
181. 181 cfg
182. 182 },
183. 183 )
184. 184}
185. 185
186. 186// Raw ledger, not the model-visible projection: these tests are about the
187. 187// repair loop's mechanics (how many drafts were tried, what a nudge said),
188. 188// which a closed turn's full record deliberately no longer preserves
189. 189// (docs/design/68-context-engine.md §10) — a rejected draft is never
190. 190// projected, and every step collapses into one trace+narration message.
191. 191fn assistant_texts(agent: &Agent) -> Vec<String> {
192. 192 futures::executor::block_on(async {
193. 193 agent
194. 194 .session
195. 195 .lock()
196. 196 .await
197. 197 .message_chain()
198. 198 .iter()
199. 199 .filter(|(_, m)| m.role == vak_llm::types::Role::Assistant)
200. 200 .flat_map(|(_, m)| m.content.iter())
201. 201 .filter_map(|b| match b {
202. 202 ContentBlock::Text { text } => Some(text.clone()),
203. 203 _ => None,
204. 204 })
205. 205 .collect()
206. 206 })
207. 207}
208. 208
209. 209#[tokio::test]
210. 210async fn ungrounded_answer_after_search_gets_one_repair_turn() {
211. 211 let dir = tempdir().unwrap();
212. 212 let mut agent = build_agent(
213. 213 &dir,
214. 214 "grounding-repair",
215. 215 vec![
216. 216 search_call("s1"),
217. 217 // First answer ignores the search result entirely — the exact
218. 218 // observed bug shape.
219. 219 text_msg("I have already provided a summary. Key themes: politics, business."),
220. 220 // After the [grounding-check] nudge, the model does the right thing.
221. 221 text_msg(&fenced_research_answer()),
222. 222 ],
223. 223 )
224. 224 .await;
225. 225
226. 226 let outcome = agent
227. 227 .run(
228. 228 "what are top news in india right now",
229. 229 &Default::default(),
230. 230 CancellationToken::new(),
231. 231 mpsc::channel(64).0,
232. 232 )
233. 233 .await;
234. 234 assert!(
235. 235 matches!(outcome, TurnOutcome::Completed { .. }),
236. 236 "got {outcome:?}"
237. 237 );
238. 238
239. 239 let texts = assistant_texts(&agent);
240. 240 assert!(
241. 241 texts.iter().any(|t| t.contains("\"semantic_type\"")),
242. 242 "final session state must contain a grounded card, got: {texts:?}"
243. 243 );
244. 244
245. 245 // The repair nudge itself must be visible in the session (for
246. 246 // diagnostics/audit), and must name the tool that was ignored.
247. 247 let user_texts: Vec<String> = futures::executor::block_on(async {
248. 248 agent
249. 249 .session
250. 250 .lock()
251. 251 .await
252. 252 .message_chain()
253. 253 .iter()
254. 254 .filter(|(_, m)| m.role == vak_llm::types::Role::User)
255. 255 .flat_map(|(_, m)| m.content.iter())
256. 256 .filter_map(|b| match b {
257. 257 ContentBlock::Text { text } => Some(text.clone()),
258. 258 _ => None,
259. 259 })
260. 260 .collect()
261. 261 });
262. 262 assert!(
263. 263 user_texts.iter().any(|t| t.contains("[grounding-check]")
264. 264 && t.contains("search")
265. 265 && t.contains("what are top news in india right now")),
266. 266 "expected a grounding-check nudge naming the ignored tool and admitted target, got: {user_texts:?}"
267. 267 );
268. 268}
269. 269
270. 270#[tokio::test]
271. 271async fn repair_is_bounded_to_one_attempt() {
272. 272 let dir = tempdir().unwrap();
273. 273 let mut agent = build_agent(
274. 274 &dir,
275. 275 "grounding-bounded",
276. 276 vec![
277. 277 search_call("s1"),
278. 278 text_msg("still vague, still no fence, first ignore"),
279. 279 // Model ignores the nudge too — must NOT loop forever; the run
280. 280 // completes (with the still-bad answer) rather than repeating.
281. 281 text_msg("still vague, still no fence, second ignore"),
282. 282 ],
283. 283 )
284. 284 .await;
285. 285
286. 286 let outcome = agent
287. 287 .run(
288. 288 "what are top news in india right now",
289. 289 &Default::default(),
290. 290 CancellationToken::new(),
291. 291 mpsc::channel(64).0,
292. 292 )
293. 293 .await;
294. 294 assert!(
295. 295 matches!(outcome, TurnOutcome::Completed { .. }),
296. 296 "got {outcome:?}"
297. 297 );
298. 298
299. 299 let texts = assistant_texts(&agent);
300. 300 assert_eq!(
301. 301 texts.iter().filter(|t| t.contains("still vague")).count(),
302. 302 2,
303. 303 "exactly one repair retry — not zero, not a loop: {texts:?}"
304. 304 );
305. 305}
306. 306
307. 307#[tokio::test]
308. 308async fn honest_no_data_admission_is_not_flagged() {
309. 309 let dir = tempdir().unwrap();
310. 310 let mut agent = build_agent(
311. 311 &dir,
312. 312 "grounding-honest",
313. 313 vec![
314. 314 search_call("s1"),
315. 315 text_msg(
316. 316 "I couldn't find current information on that from the search results — \
317. 317 the results didn't cover today's headlines.",
318. 318 ),
319. 319 ],
320. 320 )
321. 321 .await;
322. 322
323. 323 let outcome = agent
324. 324 .run(
325. 325 "what are top news in india right now",
326. 326 &Default::default(),
327. 327 CancellationToken::new(),
328. 328 mpsc::channel(64).0,
329. 329 )
330. 330 .await;
331. 331 assert!(
332. 332 matches!(outcome, TurnOutcome::Completed { .. }),
333. 333 "got {outcome:?}"
334. 334 );
335. 335
336. 336 let texts = assistant_texts(&agent);
337. 337 assert_eq!(
338. 338 texts.iter().filter(|t| t.contains("couldn't find")).count(),
339. 339 1,
340. 340 "an honest admission of no data must NOT trigger a repair retry: {texts:?}"
341. 341 );
342. 342}
343. 343
344. 344#[tokio::test]
345. 345async fn grounded_first_answer_needs_no_repair() {
346. 346 let dir = tempdir().unwrap();
347. 347 let mut agent = build_agent(
348. 348 &dir,
349. 349 "grounding-clean",
350. 350 vec![search_call("s1"), text_msg(&fenced_research_answer())],
351. 351 )
352. 352 .await;
353. 353
354. 354 let outcome = agent
355. 355 .run(
356. 356 "what are top news in india right now",
357. 357 &Default::default(),
358. 358 CancellationToken::new(),
359. 359 mpsc::channel(64).0,
360. 360 )
361. 361 .await;
362. 362 assert!(
363. 363 matches!(outcome, TurnOutcome::Completed { .. }),
364. 364 "got {outcome:?}"
365. 365 );
366. 366
367. 367 let texts = assistant_texts(&agent);
368. 368 assert_eq!(
369. 369 texts.iter().filter(|t| t.contains("semantic_type")).count(),
370. 370 1,
371. 371 "a correctly-grounded first answer must not be re-run: {texts:?}"
372. 372 );
373. 373}
374. 374
375. 375#[tokio::test]
376. 376async fn session_search_of_the_users_own_notes_is_not_flagged() {
377. 377 // Regression: `session_search` (searching the user's OWN session
378. 378 // history/notes, not external data) must never trip the grounding
379. 379 // check just because its name contains "search". This exact false
380. 380 // positive broke a real test (crates/vak-core/tests/learning_loop.rs
381. 381 // `remember_propose_recall_promote_loop`) before the fix.
382. 382 struct SessionSearchTool;
383. 383 #[async_trait]
384. 384 impl Tool for SessionSearchTool {
385. 385 fn name(&self) -> &str {
386. 386 "session_search"
387. 387 }
388. 388 fn description(&self) -> &str {
389. 389 "search the user's own session/notes history"
390. 390 }
391. 391 fn schema(&self) -> Value {
392. 392 serde_json::json!({"type": "object"})
393. 393 }
394. 394 async fn execute(&self, _args: &Value, _ctx: &ToolContext) -> ToolOutput {
395. 395 ToolOutput::ok(
396. 396 "1 hit(s): deploy-rollbacks — the deploy script must pause before rollback windows",
397. 397 )
398. 398 }
399. 399 }
400. 400
401. 401 let dir = tempdir().unwrap();
402. 402 let header = SessionHeader {
403. 403 agent: None,
404. 404 session_id: "grounding-session-search".into(),
405. 405 created_at: chrono::Utc::now(),
406. 406 cwd: dir.path().to_path_buf(),
407. 407 parent_session_id: None,
408. 408 contract_id: None,
409. 409 work_item_id: None,
410. 410 conversation: None,
411. 411 contract: FrozenContract {
412. 412 app_version: "0".into(),
413. 413 provider: "scripted".into(),
414. 414 model: "test-model".into(),
415. 415 route_ladder: Vec::new(),
416. 416 route_objective: String::new(),
417. 417 route_annotations: Vec::new(),
418. 418 system_prompt: "sys".into(),
419. 419 permission_mode: "full-access".into(),
420. 420 capabilities: Vec::new(),
421. 421 prompt_layers: Vec::new(),
422. 422 },
423. 423 };
424. 424 let home = dir.path().join("home");
425. 425 std::fs::create_dir_all(&home).unwrap();
426. 426 let log = SessionLog::create(
427. 427 SessionPath::new_session_file(&home, dir.path(), "grounding-session-search"),
428. 428 header,
429. 429 )
430. 430 .unwrap();
431. 431
432. 432 let mut agent = Agent::new(
433. 433 Arc::new(Scripted {
434. 434 responses: Mutex::new(VecDeque::from(vec![
435. 435 AssistantMessage {
436. 436 content: vec![ContentBlock::ToolUse {
437. 437 id: "s1".into(),
438. 438 name: "session_search".into(),
439. 439 input: serde_json::json!({"query": "deploy rollback"}),
440. 440 }],
441. 441 stop_reason: StopReason::ToolUse,
442. 442 usage: Usage::default(),
443. 443 model: "test-model".into(),
444. 444 response_id: None,
445. 445 },
446. 446 text_msg("Recalled from memory."),
447. 447 ])),
448. 448 }),
449. 449 log,
450. 450 {
451. 451 let mut cfg = AgentConfig::new("sys");
452. 452 cfg.model = "test-model".into();
453. 453 cfg.tools = vec![Arc::new(SessionSearchTool)];
454. 454 cfg.retrieval_check = Some(Arc::new(|name, _| name == "search"));
455. 455 cfg.mode = vak_permission::Mode::FullAccess;
456. 456 cfg.permission = Some(Arc::new(PermissionEngine::default()));
457. 457 cfg
458. 458 },
459. 459 );
460. 460
461. 461 let outcome = agent
462. 462 .run(
463. 463 "what do we know about deploys?",
464. 464 &Default::default(),
465. 465 CancellationToken::new(),
466. 466 mpsc::channel(64).0,
467. 467 )
468. 468 .await;
469. 469 assert!(
470. 470 matches!(outcome, TurnOutcome::Completed { .. }),
471. 471 "session_search must not trigger a grounding repair (and stall on an exhausted script): got {outcome:?}"
472. 472 );
473. 473
474. 474 let texts = assistant_texts(&agent);
475. 475 assert_eq!(
476. 476 texts
477. 477 .iter()
478. 478 .filter(|t| t.contains("Recalled from memory"))
479. 479 .count(),
480. 480 1,
481. 481 "no repair retry should have fired: {texts:?}"
482. 482 );
483. 483}
484. 484
485. 485#[tokio::test]
486. 486async fn ungrounded_prose_unrelated_to_any_tool_call_is_not_flagged() {
487. 487 // No search/retrieval tool ran this turn at all — an ordinary prose
488. 488 // answer must never be forced through the grounding check.
489. 489 let dir = tempdir().unwrap();
490. 490 let mut agent = build_agent(
491. 491 &dir,
492. 492 "grounding-no-tool",
493. 493 vec![text_msg("The capital of France is Paris.")],
494. 494 )
495. 495 .await;
496. 496
497. 497 let outcome = agent
498. 498 .run(
499. 499 "what is the capital of france",
500. 500 &Default::default(),
501. 501 CancellationToken::new(),
502. 502 mpsc::channel(64).0,
503. 503 )
504. 504 .await;
505. 505 assert!(
506. 506 matches!(outcome, TurnOutcome::Completed { .. }),
507. 507 "got {outcome:?}"
508. 508 );
509. 509
510. 510 let texts = assistant_texts(&agent);
511. 511 assert_eq!(
512. 512 texts
513. 513 .iter()
514. 514 .filter(|t| t.contains("capital of France"))
515. 515 .count(),
516. 516 1,
517. 517 "no tool call happened, so no grounding retry should fire: {texts:?}"
518. 518 );
519. 519}
520. 520
521. 521/// The agent never decides what is retrieval. A tool literally named `search`
522. 522/// whose results are full of URLs — exactly what the old keyword-and-URL rule
523. 523/// flagged — is not grounded on when the capability layer does not declare it
524. 524/// as reaching outside information.
525. 525#[tokio::test]
526. 526async fn a_tool_is_retrieval_only_if_the_check_says_so_whatever_its_name_or_output() {
527. 527 let dir = tempdir().unwrap();
528. 528 let mut agent = build_agent_with_check(
529. 529 &dir,
530. 530 "grounding-name-irrelevant",
531. 531 vec![
532. 532 search_call("s1"),
533. 533 text_msg("Here is a summary with no citations at all."),
534. 534 ],
535. 535 Some(Arc::new(|_: &str, _: &Value| false)),
536. 536 )
537. 537 .await;
538. 538 let outcome = agent
539. 539 .run(
540. 540 "look something up",
541. 541 &Default::default(),
542. 542 CancellationToken::new(),
543. 543 mpsc::channel(64).0,
544. 544 )
545. 545 .await;
546. 546 assert!(
547. 547 matches!(outcome, TurnOutcome::Completed { .. }),
548. 548 "got {outcome:?}"
549. 549 );
550. 550 let users: Vec<String> = futures::executor::block_on(async {
551. 551 agent
552. 552 .session
553. 553 .lock()
554. 554 .await
555. 555 .derive_messages()
556. 556 .iter()
557. 557 .filter(|m| m.role == vak_llm::types::Role::User)
558. 558 .flat_map(|m| m.content.iter())
559. 559 .filter_map(|b| match b {
560. 560 ContentBlock::Text { text } => Some(text.clone()),
561. 561 _ => None,
562. 562 })
563. 563 .collect()
564. 564 });
565. 565 assert!(
566. 566 !users.iter().any(|t| t.contains("[grounding-check]")),
567. 567 "an undeclared tool must never trigger a grounding nudge: {users:?}"
568. 568 );
569. 569}
570. 570
571. 571/// With no check configured at all the grounding check is inert rather than
572. 572/// falling back to a guess.
573. 573#[tokio::test]
574. 574async fn with_no_retrieval_check_the_grounding_check_is_inert() {
575. 575 let dir = tempdir().unwrap();
576. 576 let mut agent = build_agent_with_check(
577. 577 &dir,
578. 578 "grounding-inert",
579. 579 vec![search_call("s1"), text_msg("Uncited summary.")],
580. 580 None,
581. 581 )
582. 582 .await;
583. 583 let outcome = agent
584. 584 .run(
585. 585 "look something up",
586. 586 &Default::default(),
587. 587 CancellationToken::new(),
588. 588 mpsc::channel(64).0,
589. 589 )
590. 590 .await;
591. 591 assert!(
592. 592 matches!(outcome, TurnOutcome::Completed { .. }),
593. 593 "got {outcome:?}"
594. 594 );
595. 595}
596. 596