batty_cli/shim/
common.rs

1//! Shared utilities used by both PTY runtime and SDK runtime.
2
3use std::collections::VecDeque;
4
5use super::protocol::{Event, ShimState};
6
7// ---------------------------------------------------------------------------
8// Constants
9// ---------------------------------------------------------------------------
10
11/// Maximum number of messages that can be queued while the agent is working.
12pub const MAX_QUEUE_DEPTH: usize = 16;
13
14/// How often to report session stats (secs).
15pub const SESSION_STATS_INTERVAL_SECS: u64 = 60;
16
17// ---------------------------------------------------------------------------
18// Message formatting
19// ---------------------------------------------------------------------------
20
21pub fn reply_target_for(sender: &str) -> &str {
22    sender
23}
24
25pub fn format_injected_message(sender: &str, body: &str) -> String {
26    let reply_to = reply_target_for(sender);
27    format!(
28        "--- Message from {sender} ---\n\
29         Reply-To: {reply_to}\n\
30         If you need to reply, use: batty send {reply_to} \"<your reply>\"\n\
31         \n\
32         {body}"
33    )
34}
35
36// ---------------------------------------------------------------------------
37// Queued message
38// ---------------------------------------------------------------------------
39
40#[derive(Debug, Clone)]
41pub struct QueuedMessage {
42    pub from: String,
43    pub body: String,
44    pub message_id: Option<String>,
45}
46
47// ---------------------------------------------------------------------------
48// Queue drain helper
49// ---------------------------------------------------------------------------
50
51/// Drain all queued messages, emitting an Error event for each one.
52/// Used when the agent enters a terminal state (Dead, ContextExhausted).
53pub fn drain_queue_errors(
54    queue: &mut VecDeque<QueuedMessage>,
55    terminal_state: ShimState,
56) -> Vec<Event> {
57    let mut events = Vec::new();
58    while let Some(msg) = queue.pop_front() {
59        events.push(Event::Error {
60            command: "SendMessage".into(),
61            reason: format!(
62                "agent entered {} state, queued message dropped{}",
63                terminal_state,
64                msg.message_id
65                    .map(|id| format!(" (id: {id})"))
66                    .unwrap_or_default(),
67            ),
68        });
69    }
70    events
71}
72
73// ---------------------------------------------------------------------------
74// Context exhaustion detection
75// ---------------------------------------------------------------------------
76
77const EXHAUSTION_PATTERNS: &[&str] = &[
78    "context window exceeded",
79    "context window is full",
80    "conversation is too long",
81    "maximum context length",
82    "context limit reached",
83    "truncated due to context limit",
84    "input exceeds the model",
85    "prompt is too long",
86];
87
88/// Check if text contains known context exhaustion phrases.
89pub fn detect_context_exhausted(text: &str) -> bool {
90    let lower = text.to_lowercase();
91    EXHAUSTION_PATTERNS.iter().any(|p| lower.contains(p))
92}
93
94// ---------------------------------------------------------------------------
95// Tests
96// ---------------------------------------------------------------------------
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101
102    #[test]
103    fn format_message_includes_sender_and_body() {
104        let msg = format_injected_message("manager", "Fix the bug");
105        assert!(msg.contains("Message from manager"));
106        assert!(msg.contains("Reply-To: manager"));
107        assert!(msg.contains("Fix the bug"));
108        assert!(msg.contains("batty send manager"));
109    }
110
111    #[test]
112    fn reply_target_is_identity() {
113        assert_eq!(reply_target_for("architect"), "architect");
114        assert_eq!(reply_target_for("eng-1-1"), "eng-1-1");
115    }
116
117    #[test]
118    fn drain_queue_errors_empties_queue() {
119        let mut queue = VecDeque::new();
120        queue.push_back(QueuedMessage {
121            from: "mgr".into(),
122            body: "task 1".into(),
123            message_id: Some("id-1".into()),
124        });
125        queue.push_back(QueuedMessage {
126            from: "mgr".into(),
127            body: "task 2".into(),
128            message_id: None,
129        });
130        let events = drain_queue_errors(&mut queue, ShimState::Dead);
131        assert_eq!(events.len(), 2);
132        assert!(queue.is_empty());
133
134        match &events[0] {
135            Event::Error { reason, .. } => {
136                assert!(reason.contains("dead"));
137                assert!(reason.contains("id-1"));
138            }
139            _ => panic!("expected Error event"),
140        }
141    }
142
143    #[test]
144    fn drain_queue_empty_is_noop() {
145        let mut queue = VecDeque::new();
146        let events = drain_queue_errors(&mut queue, ShimState::Dead);
147        assert!(events.is_empty());
148    }
149
150    #[test]
151    fn context_exhaustion_detected() {
152        assert!(detect_context_exhausted("Error: context window exceeded"));
153        assert!(detect_context_exhausted(
154            "The CONVERSATION IS TOO LONG to continue"
155        ));
156        assert!(detect_context_exhausted("maximum context length reached"));
157    }
158
159    #[test]
160    fn context_exhaustion_not_detected_for_normal_text() {
161        assert!(!detect_context_exhausted("Writing function to parse YAML"));
162        assert!(!detect_context_exhausted("context manager initialized"));
163    }
164}
batty_cli/shim/common.rs

batty_cli/shim/
common.rs
