1. 1//! Regression (live-found): gateway turns with NO SSE subscriber must not
2. 2//! be aborted mid-stream when the first event fires. The old broadcast pump
3. 3//! treated zero subscribers as fatal, dropped its mpsc receiver, and the
4. 4//! agent self-cancelled on the next send — surfacing to chat users as
5. 5//! "(aborted)" exactly when the model started a tool call.
6. 6
7. 7#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
8. 8
9. 9use std::collections::VecDeque;
10. 10use std::sync::{Arc, Mutex};
11. 11
12. 12use tokio_util::sync::CancellationToken;
13. 13
14. 14use vak_core::Core;
15. 15use vak_llm::stream;
16. 16use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, Usage};
17. 17use vak_llm::{EventStream, LlmError, Provider};
18. 18
19. 19struct Scripted {
20. 20 responses: Mutex<VecDeque<AssistantMessage>>,
21. 21}
22. 22
23. 23#[async_trait::async_trait]
24. 24impl Provider for Scripted {
25. 25 fn name(&self) -> &str {
26. 26 "scripted"
27. 27 }
28. 28
29. 29 async fn stream(
30. 30 &self,
31. 31 _request: ChatRequest,
32. 32 _cancel: CancellationToken,
33. 33 ) -> Result<EventStream, LlmError> {
34. 34 let next = self.responses.lock().unwrap().pop_front();
35. 35 let (mut sink, rx) = stream::channel(64);
36. 36 match next {
37. 37 Some(m) => {
38. 38 sink.push(stream::StreamEvent::Start { partial: m.clone() });
39. 39 sink.close_message(m).await;
40. 40 }
41. 41 None => sink.close_error(LlmError::Parse("exhausted".into())).await,
42. 42 }
43. 43 Ok(rx)
44. 44 }
45. 45}
46. 46
47. 47fn text(t: &str) -> AssistantMessage {
48. 48 AssistantMessage {
49. 49 content: vec![ContentBlock::text(t)],
50. 50 stop_reason: vak_llm::types::StopReason::EndTurn,
51. 51 usage: Usage {
52. 52 input_tokens: 7,
53. 53 output_tokens: 3,
54. 54 ..Default::default()
55. 55 },
56. 56 model: "test-model".into(),
57. 57 response_id: None,
58. 58 }
59. 59}
60. 60
61. 61fn tool_call(id: &str, name: &str, input: serde_json::Value) -> AssistantMessage {
62. 62 AssistantMessage {
63. 63 content: vec![ContentBlock::ToolUse {
64. 64 id: id.into(),
65. 65 name: name.into(),
66. 66 input,
67. 67 }],
68. 68 stop_reason: vak_llm::types::StopReason::ToolUse,
69. 69 usage: Usage::default(),
70. 70 model: "test-model".into(),
71. 71 response_id: None,
72. 72 }
73. 73}
74. 74
75. 75#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
76. 76async fn headless_tool_turn_survives_without_subscribers() {
77. 77 let provider = Arc::new(Scripted {
78. 78 responses: Mutex::new(VecDeque::from(vec![
79. 79 tool_call("t1", "glob", serde_json::json!({"pattern": "*.md"})),
80. 80 text("found some markdown files"),
81. 81 ])),
82. 82 });
83. 83
84. 84 let dir = tempfile::tempdir().unwrap();
85. 85 let cwd = dir.path().to_path_buf();
86. 86 // Hermetic against the developer's global config (e.g. reflection=true):
87. 87 // pin learning flags off for deterministic scripted flows.
88. 88 let _ = std::fs::create_dir_all(cwd.join(".vak"));
89. 89 let _ = std::fs::write(
90. 90 cwd.join(".vak/config.toml"),
91. 91 "permission_mode = \"full-access\"\n[memory]\nreflection = false\n[gateway]\nchat_allowlist_open = true\n",
92. 92 );
93. 93 vak_config::paths::isolate_home_for_tests();
94. 94 let core = Core::new_with_trust(cwd.clone(), true).unwrap();
95. 95 core.set_sessions_home(dir.path().join("home"));
96. 96 core.set_permission_mode(vak_config::PermissionMode::FullAccess);
97. 97 core.set_tool_worker_exe(std::path::PathBuf::from(env!(
98. 98 "CARGO_BIN_EXE_vak-tool-worker"
99. 99 )));
100. 100 core.set_provider_instance(provider);
101. 101 std::mem::forget(dir);
102. 102
103. 103 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
104. 104 let addr = listener.local_addr().unwrap();
105. 105 tokio::spawn(async move {
106. 106 axum::serve(listener, vak_server::gateway_router(core))
107. 107 .await
108. 108 .unwrap();
109. 109 });
110. 110 let base = format!("http://{addr}");
111. 111 let client = reqwest::Client::new();
112. 112
113. 113 // Deliberately NO /events subscriber: wait:true over inbound only.
114. 114 let res = client
115. 115 .post(format!("{base}/gateway/inbound"))
116. 116 .json(&serde_json::json!({
117. 117 "surface": "probe",
118. 118 "chat": "headless",
119. 119 "text": "What files are in this project?",
120. 120 "wait": true,
121. 121 }))
122. 122 .send()
123. 123 .await
124. 124 .unwrap();
125. 125 assert_eq!(res.status(), 200);
126. 126 let body: serde_json::Value = res.json().await.unwrap();
127. 127 assert_eq!(body["state"], "completed");
128. 128 assert_ne!(
129. 129 body["text"], "(aborted)",
130. 130 "tool event must not kill a subscriber-less run"
131. 131 );
132. 132 assert_eq!(body["text"], "found some markdown files");
133. 133
134. 134 // The tool really executed: glob result is on the ledger.
135. 135 let sid = body["session_id"].as_str().unwrap().to_string();
136. 136 let t: serde_json::Value = client
137. 137 .get(format!("{base}/sessions/{sid}/transcript"))
138. 138 .send()
139. 139 .await
140. 140 .unwrap()
141. 141 .json()
142. 142 .await
143. 143 .unwrap();
144. 144 let raw = serde_json::to_string(&t).unwrap();
145. 145 assert!(raw.contains("tool_use") || raw.contains("result"), "{raw}");
146. 146}
147. 147