xz_agent_hooks/
context.rs

1//! Context injection buffer for hosts that accumulate hook text across a turn.
2
3use crate::contract::ContextChannel;
4
5/// Accumulates context strings by channel until the host drains them.
6#[derive(Debug, Clone, Default)]
7pub struct ContextInjectionBuffer {
8    items: Vec<(ContextChannel, String)>,
9}
10
11impl ContextInjectionBuffer {
12    /// Empty buffer.
13    pub fn new() -> Self {
14        Self::default()
15    }
16
17    /// Append one item.
18    pub fn push(&mut self, channel: ContextChannel, text: impl Into<String>) {
19        let text = text.into();
20        if !text.is_empty() {
21            self.items.push((channel, text));
22        }
23    }
24
25    /// Append many items.
26    pub fn extend(&mut self, items: impl IntoIterator<Item = (ContextChannel, String)>) {
27        for (c, t) in items {
28            self.push(c, t);
29        }
30    }
31
32    /// Number of pending items.
33    pub fn len(&self) -> usize {
34        self.items.len()
35    }
36
37    /// Whether empty.
38    pub fn is_empty(&self) -> bool {
39        self.items.is_empty()
40    }
41
42    /// Drain all items for one channel (preserving relative order among them).
43    pub fn drain_channel(&mut self, channel: ContextChannel) -> Vec<String> {
44        let mut kept = Vec::new();
45        let mut out = Vec::new();
46        for (c, t) in self.items.drain(..) {
47            if c == channel {
48                out.push(t);
49            } else {
50                kept.push((c, t));
51            }
52        }
53        self.items = kept;
54        out
55    }
56
57    /// Drain every item; returns (channel, text) in order.
58    pub fn drain_all(&mut self) -> Vec<(ContextChannel, String)> {
59        std::mem::take(&mut self.items)
60    }
61
62    /// Join drained PrePrompt texts with blank lines (common host helper).
63    pub fn take_pre_prompt_joined(&mut self) -> Option<String> {
64        let parts = self.drain_channel(ContextChannel::PrePrompt);
65        if parts.is_empty() {
66            None
67        } else {
68            Some(parts.join("\n\n"))
69        }
70    }
71
72    /// Peek without draining.
73    pub fn iter(&self) -> impl Iterator<Item = &(ContextChannel, String)> {
74        self.items.iter()
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81
82    #[test]
83    fn push_extend_drain() {
84        let mut b = ContextInjectionBuffer::new();
85        assert!(b.is_empty());
86        b.push(ContextChannel::PrePrompt, "a");
87        b.extend([
88            (ContextChannel::UiNotice, "u".into()),
89            (ContextChannel::PrePrompt, "b".into()),
90        ]);
91        assert_eq!(b.len(), 3);
92        let pre = b.drain_channel(ContextChannel::PrePrompt);
93        assert_eq!(pre, vec!["a", "b"]);
94        assert_eq!(b.len(), 1);
95        let joined = {
96            b.push(ContextChannel::PrePrompt, "x");
97            b.push(ContextChannel::PrePrompt, "y");
98            b.take_pre_prompt_joined()
99        };
100        assert_eq!(joined.as_deref(), Some("x\n\ny"));
101        assert_eq!(b.drain_all().len(), 1); // ui left
102    }
103
104    #[test]
105    fn empty_join_is_none() {
106        let mut b = ContextInjectionBuffer::new();
107        assert!(b.take_pre_prompt_joined().is_none());
108    }
109}
xz_agent_hooks/context.rs

xz_agent_hooks/
context.rs
