1use std::collections::{HashMap, VecDeque};
39use std::sync::Mutex;
40
41use async_trait::async_trait;
42use chrono::{DateTime, Utc};
43
44use crate::session_task::{SessionTask, TaskWakePolicy};
45use crate::task_observer::{TaskTransition, TaskTransitionObserver};
46use crate::typed_id::SessionId;
47
48#[derive(Debug, Clone, PartialEq, Eq)]
54pub struct PendingWake {
55 pub task_id: String,
56 pub session_id: SessionId,
57 pub transition: TaskTransition,
58 pub text: String,
59 pub created_at: DateTime<Utc>,
60}
61
62pub fn wake_text_for(task: &SessionTask, transition: TaskTransition) -> Option<String> {
78 match (task.wake_policy, transition) {
79 (TaskWakePolicy::Silent, _) => None,
80 (TaskWakePolicy::OnTerminal, TaskTransition::Terminal)
81 | (TaskWakePolicy::OnActivity, TaskTransition::Terminal) => {
82 let mut parts = vec![format!(
83 "Task \"{}\" ({}) finished: {}.",
84 task.display_name, task.id, task.state
85 )];
86 if let Some(summary) = &task.summary {
87 parts.push(format!("- summary: {summary}"));
88 }
89 if let Some(result_path) = &task.result_path {
90 parts.push(format!("- result_path: {result_path}"));
91 }
92 Some(parts.join("\n"))
93 }
94 (TaskWakePolicy::OnTerminal, _) => None,
95 (TaskWakePolicy::OnActivity, TaskTransition::AwaitingInput) => {
96 let prompt = task
97 .input_request
98 .as_ref()
99 .map(|r| r.prompt.as_str())
100 .unwrap_or("Task is awaiting input.");
101 Some(format!(
102 "Task \"{}\" ({}) is awaiting input: {}",
103 task.display_name, task.id, prompt
104 ))
105 }
106 (TaskWakePolicy::OnActivity, TaskTransition::Message) => {
107 let detail = task
111 .state_detail
112 .as_deref()
113 .filter(|s| !s.trim().is_empty())
114 .or_else(|| task.progress.as_ref().and_then(|p| p.label.as_deref()))
115 .unwrap_or("structured progress update");
116 Some(format!(
117 "Task \"{}\" ({}) sent a message: {}",
118 task.display_name, task.id, detail
119 ))
120 }
121 }
122}
123
124#[derive(Default)]
130pub struct SessionWakeQueue {
131 queues: Mutex<HashMap<SessionId, VecDeque<PendingWake>>>,
132}
133
134impl SessionWakeQueue {
135 pub fn new() -> Self {
136 Self::default()
137 }
138
139 pub fn note_transition(&self, task: &SessionTask, transition: TaskTransition) -> bool {
142 let Some(text) = wake_text_for(task, transition) else {
143 return false;
144 };
145 let wake = PendingWake {
146 task_id: task.id.clone(),
147 session_id: task.session_id,
148 transition,
149 text,
150 created_at: Utc::now(),
151 };
152 self.queues
153 .lock()
154 .expect("wake queue mutex poisoned")
155 .entry(task.session_id)
156 .or_default()
157 .push_back(wake);
158 true
159 }
160
161 pub fn drain(&self, session_id: SessionId) -> Vec<PendingWake> {
166 let mut guard = self.queues.lock().expect("wake queue mutex poisoned");
167 match guard.get_mut(&session_id) {
168 Some(queue) => queue.drain(..).collect(),
169 None => Vec::new(),
170 }
171 }
172
173 pub fn pending_len(&self, session_id: SessionId) -> usize {
175 self.queues
176 .lock()
177 .expect("wake queue mutex poisoned")
178 .get(&session_id)
179 .map_or(0, VecDeque::len)
180 }
181
182 pub fn has_pending(&self, session_id: SessionId) -> bool {
184 self.pending_len(session_id) > 0
185 }
186}
187
188#[async_trait]
193impl TaskTransitionObserver for SessionWakeQueue {
194 async fn on_transition(
195 &self,
196 task: &SessionTask,
197 transition: TaskTransition,
198 ) -> anyhow::Result<()> {
199 self.note_transition(task, transition);
200 Ok(())
201 }
202}
203
204#[cfg(test)]
205mod tests {
206 use super::*;
207 use crate::session_task::{
208 CreateSessionTask, SessionTaskState, TaskInputRequest, new_session_task,
209 };
210
211 fn task_with_policy(policy: TaskWakePolicy) -> SessionTask {
212 let session_id = SessionId::new();
213 let mut task = new_session_task(
214 CreateSessionTask {
215 id: None,
216 session_id,
217 kind: "subagent".into(),
218 display_name: "Test Runner".into(),
219 spec: serde_json::Value::Null,
220 state: SessionTaskState::Running,
221 links: Default::default(),
222 wake_policy: policy,
223 },
224 Utc::now(),
225 );
226 task.state = SessionTaskState::Succeeded;
227 task
228 }
229
230 #[test]
231 fn silent_never_wakes() {
232 for transition in [
233 TaskTransition::Terminal,
234 TaskTransition::AwaitingInput,
235 TaskTransition::Message,
236 ] {
237 assert!(wake_text_for(&task_with_policy(TaskWakePolicy::Silent), transition).is_none());
238 }
239 }
240
241 #[test]
242 fn on_terminal_wakes_only_on_terminal() {
243 let task = task_with_policy(TaskWakePolicy::OnTerminal);
244 assert!(wake_text_for(&task, TaskTransition::Terminal).is_some());
245 assert!(wake_text_for(&task, TaskTransition::AwaitingInput).is_none());
246 assert!(wake_text_for(&task, TaskTransition::Message).is_none());
247 }
248
249 #[test]
250 fn on_activity_wakes_on_all_three() {
251 let mut task = task_with_policy(TaskWakePolicy::OnActivity);
252 task.state = SessionTaskState::AwaitingInput;
253 task.input_request = Some(TaskInputRequest {
254 id: "ir_1".into(),
255 prompt: "pick a branch".into(),
256 expected: None,
257 });
258 let awaiting = wake_text_for(&task, TaskTransition::AwaitingInput).expect("awaiting wakes");
259 assert!(awaiting.contains("awaiting input"));
260 assert!(awaiting.contains("pick a branch"));
261
262 task.state = SessionTaskState::Running;
263 task.state_detail = Some("iteration 4/10".into());
264 let message = wake_text_for(&task, TaskTransition::Message).expect("message wakes");
265 assert!(message.contains("iteration 4/10"));
266
267 task.state = SessionTaskState::Failed;
268 assert!(wake_text_for(&task, TaskTransition::Terminal).is_some());
269 }
270
271 #[test]
272 fn terminal_text_includes_summary_and_result_path() {
273 let mut task = task_with_policy(TaskWakePolicy::OnTerminal);
274 task.summary = Some("all tests passed".into());
275 task.result_path = Some("/.tasks/task_x/result.json".into());
276 let text = wake_text_for(&task, TaskTransition::Terminal).unwrap();
277 assert!(text.contains("finished: succeeded"));
278 assert!(text.contains("all tests passed"));
279 assert!(text.contains("/.tasks/task_x/result.json"));
280 }
281
282 #[test]
283 fn drain_is_exactly_once() {
284 let queue = SessionWakeQueue::new();
285 let task = task_with_policy(TaskWakePolicy::OnTerminal);
286 let session_id = task.session_id;
287
288 assert!(queue.note_transition(&task, TaskTransition::Terminal));
289 assert_eq!(queue.pending_len(session_id), 1);
290
291 let first = queue.drain(session_id);
292 assert_eq!(first.len(), 1, "first drain returns the wake");
293 assert_eq!(first[0].task_id, task.id);
294
295 let second = queue.drain(session_id);
296 assert!(
297 second.is_empty(),
298 "second drain returns nothing — claimed once"
299 );
300 assert_eq!(queue.pending_len(session_id), 0);
301 }
302
303 #[test]
304 fn silent_transition_enqueues_nothing() {
305 let queue = SessionWakeQueue::new();
306 let task = task_with_policy(TaskWakePolicy::Silent);
307 assert!(!queue.note_transition(&task, TaskTransition::Terminal));
308 assert_eq!(queue.pending_len(task.session_id), 0);
309 }
310
311 #[test]
312 fn queues_are_isolated_per_session() {
313 let queue = SessionWakeQueue::new();
314 let a = task_with_policy(TaskWakePolicy::OnTerminal);
315 let b = task_with_policy(TaskWakePolicy::OnTerminal);
316 queue.note_transition(&a, TaskTransition::Terminal);
317 queue.note_transition(&b, TaskTransition::Terminal);
318 assert_eq!(queue.drain(a.session_id).len(), 1);
319 assert_eq!(
320 queue.pending_len(b.session_id),
321 1,
322 "draining A leaves B intact"
323 );
324 }
325}