agentd/subagent/
replies.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//! The child-side **reply slots** for supervisor round-trips: a turn worker
3//! sends `ToolRequest`/`BudgetRequest` frames up and blocks until the
4//! control-reader thread delivers the matching `ToolResult`/`BudgetGrant` here,
5//! matched by `id`. One mutex plus one condvar guards the whole map, and a
6//! closed channel or a set cancel flag wakes every waiter with `None` so no
7//! worker can outlive the supervisor that was going to answer it.
8
9use serde_json::Value;
10use std::collections::HashMap;
11use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
12use std::sync::{Condvar, Mutex};
13use std::time::{Duration, Instant};
14
15/// A delivered reply.
16#[derive(Debug, Clone, PartialEq)]
17pub enum Reply {
18    Tool {
19        result: Value,
20        is_error: bool,
21    },
22    Budget {
23        ok: bool,
24        wait_ms: u64,
25        model: Option<String>,
26        reason: Option<String>,
27    },
28}
29
30#[derive(Default)]
31pub struct Replies {
32    slots: Mutex<HashMap<u64, Reply>>,
33    cv: Condvar,
34    closed: AtomicBool,
35    next_id: AtomicU64,
36}
37
38impl Replies {
39    pub fn new() -> Replies {
40        Replies {
41            next_id: AtomicU64::new(1),
42            ..Default::default()
43        }
44    }
45
46    /// Mint a fresh request id.
47    pub fn next_id(&self) -> u64 {
48        self.next_id.fetch_add(1, Ordering::Relaxed)
49    }
50
51    /// Deliver a reply (the control thread).
52    pub fn deliver(&self, id: u64, reply: Reply) {
53        self.slots
54            .lock()
55            .unwrap_or_else(|e| e.into_inner())
56            .insert(id, reply);
57        self.cv.notify_all();
58    }
59
60    /// The channel closed (supervisor gone): wake everyone.
61    pub fn close(&self) {
62        self.closed.store(true, Ordering::Relaxed);
63        self.cv.notify_all();
64    }
65
66    pub fn is_closed(&self) -> bool {
67        self.closed.load(Ordering::Relaxed)
68    }
69
70    /// Block until the reply for `id` arrives, the deadline passes, the channel
71    /// closes, or `cancel` is set. Cancel is not signalled through the condvar,
72    /// so the wait is capped at 100 ms per iteration to keep polling it.
73    /// Returns `None` for every non-delivery ending. An already-delivered reply
74    /// is claimed before any of those checks run, so a reply that landed before
75    /// the call wins even against an expired deadline.
76    pub fn wait(&self, id: u64, deadline: Instant, cancel: &AtomicBool) -> Option<Reply> {
77        let mut slots = self.slots.lock().unwrap_or_else(|e| e.into_inner());
78        loop {
79            if let Some(r) = slots.remove(&id) {
80                return Some(r);
81            }
82            if self.is_closed() || cancel.load(Ordering::Relaxed) {
83                return None;
84            }
85            let now = Instant::now();
86            if now >= deadline {
87                return None;
88            }
89            let wait = (deadline - now).min(Duration::from_millis(100));
90            let (guard, _) = self
91                .cv
92                .wait_timeout(slots, wait)
93                .unwrap_or_else(|e| e.into_inner());
94            slots = guard;
95        }
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102    use serde_json::json;
103    use std::sync::Arc;
104
105    #[test]
106    fn replies_are_delivered_by_id_and_waiters_wake_on_close_or_cancel() {
107        let r = Arc::new(Replies::new());
108        let id = r.next_id();
109        let r2 = r.clone();
110        let t = std::thread::spawn(move || {
111            std::thread::sleep(Duration::from_millis(30));
112            r2.deliver(
113                99,
114                Reply::Tool {
115                    result: json!("other"),
116                    is_error: false,
117                },
118            );
119            r2.deliver(
120                id,
121                Reply::Tool {
122                    result: json!({"ok": true}),
123                    is_error: false,
124                },
125            );
126        });
127        let cancel = AtomicBool::new(false);
128        let got = r
129            .wait(id, Instant::now() + Duration::from_secs(2), &cancel)
130            .unwrap();
131        assert_eq!(
132            got,
133            Reply::Tool {
134                result: json!({"ok": true}),
135                is_error: false
136            }
137        );
138        t.join().unwrap();
139        // Timeout.
140        assert!(
141            r.wait(12345, Instant::now() + Duration::from_millis(20), &cancel)
142                .is_none()
143        );
144        // Cancel wakes.
145        cancel.store(true, Ordering::Relaxed);
146        assert!(
147            r.wait(12345, Instant::now() + Duration::from_secs(5), &cancel)
148                .is_none()
149        );
150        // Close wakes.
151        let cancel2 = AtomicBool::new(false);
152        let r3 = r.clone();
153        let t = std::thread::spawn(move || {
154            std::thread::sleep(Duration::from_millis(20));
155            r3.close();
156        });
157        assert!(
158            r.wait(777, Instant::now() + Duration::from_secs(5), &cancel2)
159                .is_none()
160        );
161        t.join().unwrap();
162        // A reply that arrived before the wait is picked up.
163        r.deliver(
164            5,
165            Reply::Budget {
166                ok: true,
167                wait_ms: 0,
168                model: None,
169                reason: None,
170            },
171        );
172        assert!(matches!(
173            r.wait(5, Instant::now(), &cancel2),
174            Some(Reply::Budget { ok: true, .. })
175        ));
176    }
177}
agentd/subagent/replies.rs

agentd/subagent/
replies.rs
