xz-agent-hooks 0.1.0

//! Context injection buffer for hosts that accumulate hook text across a turn.

use crate::contract::ContextChannel;

/// Accumulates context strings by channel until the host drains them.
#[derive(Debug, Clone, Default)]
pub struct ContextInjectionBuffer {
    items: Vec<(ContextChannel, String)>,
}

impl ContextInjectionBuffer {
    /// Empty buffer.
    pub fn new() -> Self {
        Self::default()
    }

    /// Append one item.
    pub fn push(&mut self, channel: ContextChannel, text: impl Into<String>) {
        let text = text.into();
        if !text.is_empty() {
            self.items.push((channel, text));
        }
    }

    /// Append many items.
    pub fn extend(&mut self, items: impl IntoIterator<Item = (ContextChannel, String)>) {
        for (c, t) in items {
            self.push(c, t);
        }
    }

    /// Number of pending items.
    pub fn len(&self) -> usize {
        self.items.len()
    }

    /// Whether empty.
    pub fn is_empty(&self) -> bool {
        self.items.is_empty()
    }

    /// Drain all items for one channel (preserving relative order among them).
    pub fn drain_channel(&mut self, channel: ContextChannel) -> Vec<String> {
        let mut kept = Vec::new();
        let mut out = Vec::new();
        for (c, t) in self.items.drain(..) {
            if c == channel {
                out.push(t);
            } else {
                kept.push((c, t));
            }
        }
        self.items = kept;
        out
    }

    /// Drain every item; returns (channel, text) in order.
    pub fn drain_all(&mut self) -> Vec<(ContextChannel, String)> {
        std::mem::take(&mut self.items)
    }

    /// Join drained PrePrompt texts with blank lines (common host helper).
    pub fn take_pre_prompt_joined(&mut self) -> Option<String> {
        let parts = self.drain_channel(ContextChannel::PrePrompt);
        if parts.is_empty() {
            None
        } else {
            Some(parts.join("\n\n"))
        }
    }

    /// Peek without draining.
    pub fn iter(&self) -> impl Iterator<Item = &(ContextChannel, String)> {
        self.items.iter()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn push_extend_drain() {
        let mut b = ContextInjectionBuffer::new();
        assert!(b.is_empty());
        b.push(ContextChannel::PrePrompt, "a");
        b.extend([
            (ContextChannel::UiNotice, "u".into()),
            (ContextChannel::PrePrompt, "b".into()),
        ]);
        assert_eq!(b.len(), 3);
        let pre = b.drain_channel(ContextChannel::PrePrompt);
        assert_eq!(pre, vec!["a", "b"]);
        assert_eq!(b.len(), 1);
        let joined = {
            b.push(ContextChannel::PrePrompt, "x");
            b.push(ContextChannel::PrePrompt, "y");
            b.take_pre_prompt_joined()
        };
        assert_eq!(joined.as_deref(), Some("x\n\ny"));
        assert_eq!(b.drain_all().len(), 1); // ui left
    }

    #[test]
    fn empty_join_is_none() {
        let mut b = ContextInjectionBuffer::new();
        assert!(b.take_pre_prompt_joined().is_none());
    }
}