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_util::sync::CancellationToken;
7. 7
8. 8use vak_core::Core;
9. 9use vak_llm::stream;
10. 10use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, StopReason, Usage};
11. 11use vak_llm::{EventStream, LlmError, Provider};
12. 12
13. 13struct Scripted {
14. 14 responses: Mutex<VecDeque<AssistantMessage>>,
15. 15}
16. 16
17. 17#[async_trait::async_trait]
18. 18impl Provider for Scripted {
19. 19 fn name(&self) -> &str {
20. 20 "scripted"
21. 21 }
22. 22
23. 23 async fn stream(
24. 24 &self,
25. 25 _request: ChatRequest,
26. 26 _cancel: CancellationToken,
27. 27 ) -> Result<EventStream, LlmError> {
28. 28 let next = self.responses.lock().unwrap().pop_front();
29. 29 let (mut sink, rx) = stream::channel(64);
30. 30 match next {
31. 31 Some(m) => {
32. 32 sink.push(stream::StreamEvent::Start { partial: m.clone() });
33. 33 sink.close_message(m).await;
34. 34 }
35. 35 None => sink.close_error(LlmError::Parse("exhausted".into())).await,
36. 36 }
37. 37 Ok(rx)
38. 38 }
39. 39}
40. 40
41. 41fn text(t: &str) -> AssistantMessage {
42. 42 AssistantMessage {
43. 43 content: vec![ContentBlock::text(t)],
44. 44 stop_reason: StopReason::EndTurn,
45. 45 usage: Usage {
46. 46 input_tokens: 7,
47. 47 output_tokens: 3,
48. 48 ..Default::default()
49. 49 },
50. 50 model: "test-model".into(),
51. 51 response_id: None,
52. 52 }
53. 53}
54. 54
55. 55#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
56. 56async fn mcp_servers_get_put_roundtrip_and_persist() {
57. 57 let dir = tempfile::tempdir().unwrap();
58. 58 let cwd = dir.path().to_path_buf();
59. 59 let project = cwd.join(".vak");
60. 60 std::fs::create_dir_all(&project).unwrap();
61. 61 // Pre-existing config with a comment-bearing key we must not destroy.
62. 62 std::fs::write(
63. 63 project.join("config.toml"),
64. 64 "model = \"claude-sonnet-4-5\"\nmax_turns = 12\n",
65. 65 )
66. 66 .unwrap();
67. 67
68. 68 vak_config::paths::isolate_home_for_tests();
69. 69 let core = Core::new_with_trust(cwd.clone(), true).expect("core");
70. 70 core.set_sessions_home(dir.path().join("home"));
71. 71 core.set_provider_instance(Arc::new(Scripted {
72. 72 responses: Mutex::new(VecDeque::from(vec![text("ok")])),
73. 73 }));
74. 74 std::mem::forget(dir);
75. 75
76. 76 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
77. 77 let addr = listener.local_addr().unwrap();
78. 78 let app = vak_server::router(core);
79. 79 tokio::spawn(async move {
80. 80 axum::serve(listener, app).await.unwrap();
81. 81 });
82. 82 let base = format!("http://{addr}");
83. 83 let client = reqwest::Client::new();
84. 84
85. 85 // Empty at first (this core has no servers of its own; a developer
86. 86 // machine's global config may inject some, so only assert shape).
87. 87 let res = client
88. 88 .get(format!("{base}/config/mcp"))
89. 89 .send()
90. 90 .await
91. 91 .unwrap();
92. 92 assert_eq!(res.status(), 200);
93. 93 let body: serde_json::Value = res.json().await.unwrap();
94. 94 assert!(body["servers"].is_object());
95. 95
96. 96 // PUT replaces the whole RUNNING table (global + project) and persists
97. 97 // the project-layer file.
98. 98
99. 99 // PUT two servers (one with env + network).
100. 100 let res = client
101. 101 .put(format!("{base}/config/mcp"))
102. 102 .json(&serde_json::json!({
103. 103 "servers": {
104. 104 "context7": {"command": "npx", "args": ["-y", "@context7/mcp"]},
105. 105 "fs": {"command": "./bin/fsd", "env": {"TOKEN": "${FS_TOKEN}"}, "network": true}
106. 106 }
107. 107 }))
108. 108 .send()
109. 109 .await
110. 110 .unwrap();
111. 111 assert_eq!(res.status(), 200, "{res:?}");
112. 112
113. 113 // GET reflects the hot-applied table immediately — exactly what was
114. 114 // PUT, nothing inherited.
115. 115 let res = client
116. 116 .get(format!("{base}/config/mcp"))
117. 117 .send()
118. 118 .await
119. 119 .unwrap();
120. 120 let body: serde_json::Value = res.json().await.unwrap();
121. 121 assert_eq!(
122. 122 body["servers"],
123. 123 serde_json::json!({
124. 124 "context7": {"command": "npx", "args": ["-y", "@context7/mcp"], "env": {}, "network": false},
125. 125 "fs": {"command": "./bin/fsd", "args": [], "env": {"TOKEN": "${FS_TOKEN}"}, "network": true},
126. 126 })
127. 127 );
128. 128
129. 129 // Persisted into the project config without destroying other keys.
130. 130 let raw = std::fs::read_to_string(cwd.join(".vak/config.toml")).unwrap();
131. 131 let parsed: toml::Value = toml::from_str(&raw).unwrap();
132. 132 assert_eq!(
133. 133 parsed["model"],
134. 134 toml::Value::from("claude-sonnet-4-5"),
135. 135 "existing keys survive"
136. 136 );
137. 137 assert_eq!(parsed["max_turns"], toml::Value::Integer(12));
138. 138 assert_eq!(
139. 139 parsed["mcp"]["servers"]["context7"]["command"],
140. 140 toml::Value::from("npx")
141. 141 );
142. 142
143. 143 // Invalid names are rejected and change nothing.
144. 144 let res = client
145. 145 .put(format!("{base}/config/mcp"))
146. 146 .json(&serde_json::json!({
147. 147 "servers": {"bad name!": {"command": "x"}}
148. 148 }))
149. 149 .send()
150. 150 .await
151. 151 .unwrap();
152. 152 assert_eq!(res.status(), 400);
153. 153 let res = client
154. 154 .put(format!("{base}/config/mcp"))
155. 155 .json(&serde_json::json!({"servers": {"empty-cmd": {"command": " "}}}))
156. 156 .send()
157. 157 .await
158. 158 .unwrap();
159. 159 assert_eq!(res.status(), 400);
160. 160}
161. 161
162. 162/// `PUT /config/mcp` and `PATCH /config` arriving at the same moment each
163. 163/// rewrite the project's `.vak/config.toml`. Both must succeed and both
164. 164/// changes must land, round after round, in a file that still parses.
165. 165#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
166. 166async fn put_mcp_and_patch_config_at_the_same_moment_both_land() {
167. 167 use axum::body::Body;
168. 168 use axum::http::{Request, StatusCode};
169. 169 use tower::ServiceExt;
170. 170
171. 171 let dir = tempfile::tempdir().unwrap();
172. 172 let cwd = dir.path().join("workspace");
173. 173 std::fs::create_dir_all(cwd.join(".vak")).unwrap();
174. 174 vak_config::paths::isolate_home_for_tests();
175. 175 let core = Core::new_with_trust(cwd.clone(), true).expect("core");
176. 176 core.set_sessions_home(dir.path().join("home"));
177. 177 let app = vak_server::router(core);
178. 178 let send = |method: &'static str, uri: &'static str, body: serde_json::Value| {
179. 179 let request = Request::builder()
180. 180 .method(method)
181. 181 .uri(uri)
182. 182 .header("content-type", "application/json")
183. 183 .body(Body::from(body.to_string()))
184. 184 .unwrap();
185. 185 tokio::spawn(app.clone().oneshot(request))
186. 186 };
187. 187
188. 188 for round in 1..=30_u32 {
189. 189 let server = format!("server-{round}");
190. 190 let put = send(
191. 191 "PUT",
192. 192 "/config/mcp",
193. 193 serde_json::json!({ "servers": { server.clone(): { "command": "npx" } } }),
194. 194 );
195. 195 let patch = send(
196. 196 "PATCH",
197. 197 "/config",
198. 198 serde_json::json!({ "max_turns": round }),
199. 199 );
200. 200 let put = put.await.unwrap().unwrap().status();
201. 201 let patch = patch.await.unwrap().unwrap().status();
202. 202 assert_eq!(
203. 203 (put, patch),
204. 204 (StatusCode::OK, StatusCode::OK),
205. 205 "round {round}"
206. 206 );
207. 207
208. 208 let text = std::fs::read_to_string(cwd.join(".vak/config.toml")).unwrap();
209. 209 let document: toml::Table = toml::from_str(&text).expect("the file parses");
210. 210 let servers = document["mcp"]["servers"].as_table().unwrap();
211. 211 assert!(
212. 212 servers.contains_key(&server),
213. 213 "round {round}: the MCP change was lost: {text}"
214. 214 );
215. 215 assert_eq!(
216. 216 document["max_turns"].as_integer(),
217. 217 Some(i64::from(round)),
218. 218 "round {round}: the PATCH was lost: {text}"
219. 219 );
220. 220 }
221. 221}
222. 222