agentd/subagent/
replies.rs

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

agentd/subagent/
replies.rs
