1. 1#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
2. 2
3. 3use std::collections::VecDeque;
4. 4use std::path::Path;
5. 5use std::sync::{Arc, Mutex};
6. 6
7. 7use tokio_util::sync::CancellationToken;
8. 8
9. 9use vak_core::Core;
10. 10use vak_llm::stream;
11. 11use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, Usage};
12. 12use vak_llm::{EventStream, LlmError, Provider};
13. 13use vak_session::types::{FrozenContract, MessageRecord, SessionHeader};
14. 14use vak_session::{SessionLog, SessionPath};
15. 15
16. 16struct Scripted {
17. 17 responses: Mutex<VecDeque<AssistantMessage>>,
18. 18}
19. 19
20. 20#[async_trait::async_trait]
21. 21impl Provider for Scripted {
22. 22 fn name(&self) -> &str {
23. 23 "scripted"
24. 24 }
25. 25
26. 26 async fn stream(
27. 27 &self,
28. 28 _request: ChatRequest,
29. 29 _cancel: CancellationToken,
30. 30 ) -> Result<EventStream, LlmError> {
31. 31 let next = self.responses.lock().unwrap().pop_front();
32. 32 let (mut sink, rx) = stream::channel(64);
33. 33 match next {
34. 34 Some(m) => {
35. 35 sink.push(stream::StreamEvent::Start { partial: m.clone() });
36. 36 sink.close_message(m).await;
37. 37 }
38. 38 None => sink.close_error(LlmError::Parse("exhausted".into())).await,
39. 39 }
40. 40 Ok(rx)
41. 41 }
42. 42}
43. 43
44. 44fn text(t: &str) -> AssistantMessage {
45. 45 AssistantMessage {
46. 46 content: vec![ContentBlock::text(t)],
47. 47 stop_reason: vak_llm::types::StopReason::EndTurn,
48. 48 usage: Usage {
49. 49 input_tokens: 7,
50. 50 output_tokens: 3,
51. 51 ..Default::default()
52. 52 },
53. 53 model: "test-model".into(),
54. 54 response_id: None,
55. 55 }
56. 56}
57. 57
58. 58fn tool_call(id: &str, name: &str, input: serde_json::Value) -> AssistantMessage {
59. 59 AssistantMessage {
60. 60 content: vec![ContentBlock::ToolUse {
61. 61 id: id.into(),
62. 62 name: name.into(),
63. 63 input,
64. 64 }],
65. 65 stop_reason: vak_llm::types::StopReason::ToolUse,
66. 66 usage: Usage::default(),
67. 67 model: "test-model".into(),
68. 68 response_id: None,
69. 69 }
70. 70}
71. 71
72. 72fn header_for(id: &str, cwd: &Path) -> SessionHeader {
73. 73 SessionHeader {
74. 74 agent: Some(vak_session::types::AgentIdentity {
75. 75 id: "vak".into(),
76. 76 revision: 1,
77. 77 name: "Vak".into(),
78. 78 character: "vak".into(),
79. 79 personality: String::new(),
80. 80 animation: "subtle".into(),
81. 81 voice: "default".into(),
82. 82 behaviour: String::new(),
83. 83 responsibilities: String::new(),
84. 84 instructions: String::new(),
85. 85 }),
86. 86 session_id: id.to_string(),
87. 87 created_at: chrono::Utc::now(),
88. 88 cwd: cwd.to_path_buf(),
89. 89 parent_session_id: None,
90. 90 contract_id: None,
91. 91 work_item_id: None,
92. 92 conversation: Some(vak_session::types::ConversationContext::local(
93. 93 format!("conversation:{id}"),
94. 94 "local",
95. 95 )),
96. 96 contract: FrozenContract {
97. 97 app_version: "test".into(),
98. 98 provider: "scripted".into(),
99. 99 model: "m".into(),
100. 100 route_ladder: Vec::new(),
101. 101 route_objective: String::new(),
102. 102 route_annotations: Vec::new(),
103. 103 system_prompt: String::new(),
104. 104 permission_mode: "workspace-write".into(),
105. 105 capabilities: Vec::new(),
106. 106 prompt_layers: Vec::new(),
107. 107 },
108. 108 }
109. 109}
110. 110
111. 111#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
112. 112async fn session_search_tool_is_available_and_logged() {
113. 113 let dir = tempfile::tempdir().unwrap();
114. 114 let home = dir.path().join("home");
115. 115 let cwd = dir.path().to_path_buf();
116. 116
117. 117 // A past session with retrievable knowledge.
118. 118 let past_path = SessionPath::new_session_file(&home, &cwd, "11111111-past");
119. 119 let mut past = SessionLog::create(past_path, header_for("11111111-past", &cwd)).unwrap();
120. 120 past.append_message(MessageRecord {
121. 121 message: vak_llm::Message {
122. 122 role: vak_llm::Role::User,
123. 123 content: vec![ContentBlock::text(
124. 124 "we decided the deploy script must pause before rollback windows",
125. 125 )],
126. 126 },
127. 127 meta: None,
128. 128 })
129. 129 .unwrap();
130. 130 drop(past); // release the ledger lock
131. 131
132. 132 // Live run: the scripted model reaches for memory, then answers.
133. 133 vak_config::paths::isolate_home_for_tests();
134. 134 let core = Core::new(cwd.clone()).unwrap();
135. 135 core.set_sessions_home(home.clone());
136. 136 core.set_permission_mode(vak_config::PermissionMode::FullAccess);
137. 137 core.set_provider_instance(Arc::new(Scripted {
138. 138 responses: Mutex::new(VecDeque::from(vec![
139. 139 tool_call(
140. 140 "t1",
141. 141 "session_search",
142. 142 serde_json::json!({"query": "deploy script"}),
143. 143 ),
144. 144 text("Found it: the deploy script pauses before rollbacks."),
145. 145 ])),
146. 146 }));
147. 147 std::mem::forget(dir);
148. 148
149. 149 // Direct probe of the search function on identical inputs.
150. 150 let probe = vak_session::search(&home, &cwd, "deploy script", 8, &Default::default()).unwrap();
151. 151 eprintln!("PROBE hits={}", probe.len());
152. 152 for h in &probe {
153. 153 eprintln!(
154. 154 " -> {} score={} snippet={}",
155. 155 h.session_id, h.score, h.snippet
156. 156 );
157. 157 }
158. 158
159. 159 let live = core.start_session().await.unwrap();
160. 160 let live_id = live.header().unwrap().session_id.clone();
161. 161 let (events_tx, _events_rx) = tokio::sync::mpsc::channel(256);
162. 162 let result = core
163. 163 .run_turn_with(
164. 164 live,
165. 165 "what did we decide about deploys?",
166. 166 CancellationToken::new(),
167. 167 None,
168. 168 None,
169. 169 None,
170. 170 events_tx,
171. 171 )
172. 172 .await
173. 173 .unwrap();
174. 174 assert!(matches!(result.0, vak_agent::TurnOutcome::Completed { .. }));
175. 175 let log = result.1;
176. 176
177. 177 // Invariant 1: the search output must be reconstructable from the
178. 178 // ledger — it lives in the ToolResult block's content.
179. 179 let chain = log.message_chain();
180. 180 let mut saw_search_output = false;
181. 181 for (_, m) in &chain {
182. 182 for b in &m.content {
183. 183 if let ContentBlock::ToolResult {
184. 184 content, is_error, ..
185. 185 } = b
186. 186 {
187. 187 eprintln!("TOOL_RESULT[{is_error}]: {content}");
188. 188 assert!(!is_error, "search must not error: {content}");
189. 189 if content.contains("hit(s)") {
190. 190 saw_search_output = true;
191. 191 assert!(
192. 192 content.contains("11111111-past"),
193. 193 "snippet cites the source session: {content}"
194. 194 );
195. 195 assert!(
196. 196 !content.contains(&live_id),
197. 197 "current session must be excluded from results"
198. 198 );
199. 199 }
200. 200 }
201. 201 }
202. 202 }
203. 203 assert!(
204. 204 saw_search_output,
205. 205 "tool result with hits must be logged on the chain"
206. 206 );
207. 207
208. 208 // And the model's answer reflects what memory returned.
209. 209 let answered = chain
210. 210 .iter()
211. 211 .any(|(_, m)| m.text_content().contains("pauses before rollbacks"));
212. 212 assert!(answered, "final answer on the chain");
213. 213
214. 214 // The tool appears in the frozen contract.
215. 215 let capabilities = &log.header().unwrap().contract.capabilities;
216. 216 assert!(
217. 217 capabilities.iter().any(|capability| {
218. 218 capability.kind == vak_session::CapabilityKind::Tool
219. 219 && capability.name == "session_search"
220. 220 }),
221. 221 "session_search in frozen contract: {capabilities:?}"
222. 222 );
223. 223}
224. 224