use std::sync::{Arc, Mutex, MutexGuard};
const MAX_PENDING_STEERING_MESSAGES: usize = 16;
pub(crate) const STEERING_HEADER: &str = "Steering update from user while current run was active:";
#[derive(Clone, Debug)]
pub(crate) struct AgentSteering {
inner: Arc<Mutex<SteeringQueue>>,
}
#[derive(Debug, Default)]
struct SteeringQueue {
messages: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct SteeringQueued {
pub(crate) position: usize,
pub(crate) pending_count: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct SteeringBatch {
pub(crate) text: String,
pub(crate) count: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum SteeringRejected {
Empty,
Full { capacity: usize },
}
impl AgentSteering {
pub(crate) fn new() -> Self {
Self {
inner: Arc::new(Mutex::new(SteeringQueue::default())),
}
}
pub(crate) fn try_enqueue(&self, text: String) -> Result<SteeringQueued, SteeringRejected> {
let text = text.trim().to_string();
if text.is_empty() {
return Err(SteeringRejected::Empty);
}
let mut queue = self.lock_queue();
if queue.messages.len() >= MAX_PENDING_STEERING_MESSAGES {
return Err(SteeringRejected::Full {
capacity: MAX_PENDING_STEERING_MESSAGES,
});
}
let position = queue.messages.len();
queue.messages.push(text);
Ok(SteeringQueued {
position,
pending_count: queue.messages.len(),
})
}
pub(crate) fn observe_collapsed(&self) -> Option<SteeringBatch> {
let queue = self.lock_queue();
(!queue.messages.is_empty()).then(|| SteeringBatch {
text: collapse_steering_messages(&queue.messages),
count: queue.messages.len(),
})
}
pub(crate) fn acknowledge(&self, count: usize) {
let mut queue = self.lock_queue();
let acknowledged = count.min(queue.messages.len());
queue.messages.drain(..acknowledged);
}
pub(crate) fn acknowledge_collapsed(&self, text: &str) -> bool {
let mut queue = self.lock_queue();
let Some(count) = (1..=queue.messages.len())
.find(|count| collapse_steering_messages(&queue.messages[..*count]) == text)
else {
return false;
};
queue.messages.drain(..count);
true
}
#[cfg(test)]
pub(crate) fn drain_collapsed(&self) -> Option<String> {
let batch = self.observe_collapsed()?;
self.acknowledge(batch.count);
Some(batch.text)
}
pub(crate) fn pending_count(&self) -> usize {
self.lock_queue().messages.len()
}
pub(crate) fn clear(&self) -> usize {
let mut queue = self.lock_queue();
let count = queue.messages.len();
queue.messages.clear();
count
}
fn lock_queue(&self) -> MutexGuard<'_, SteeringQueue> {
self.inner
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
}
impl Default for AgentSteering {
fn default() -> Self {
Self::new()
}
}
fn collapse_steering_messages(messages: &[String]) -> String {
if messages.len() == 1 {
return format!("{STEERING_HEADER}\n\n{}", messages[0]);
}
let numbered = messages
.iter()
.enumerate()
.map(|(index, message)| format!("{}. {message}", index + 1))
.collect::<Vec<_>>()
.join("\n\n");
format!("{STEERING_HEADER}\n\n{numbered}")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn enqueue_single_message_drain_returns_header_and_message() {
let steering = AgentSteering::new();
assert_eq!(
steering.try_enqueue("keep going".to_string()),
Ok(SteeringQueued {
position: 0,
pending_count: 1,
})
);
assert_eq!(
steering.drain_collapsed(),
Some(
"Steering update from user while current run was active:\n\nkeep going".to_string()
)
);
}
#[test]
fn enqueue_multiple_messages_drain_returns_numbered_format() {
let steering = AgentSteering::new();
assert_eq!(
steering.try_enqueue("first".to_string()),
Ok(SteeringQueued {
position: 0,
pending_count: 1,
})
);
assert_eq!(
steering.try_enqueue("second".to_string()),
Ok(SteeringQueued {
position: 1,
pending_count: 2,
})
);
assert_eq!(
steering.try_enqueue("third".to_string()),
Ok(SteeringQueued {
position: 2,
pending_count: 3,
})
);
assert_eq!(
steering.drain_collapsed(),
Some(
"Steering update from user while current run was active:\n\n1. first\n\n2. second\n\n3. third"
.to_string()
)
);
}
#[test]
fn enqueue_empty_or_whitespace_only_text_rejects_empty() {
let steering = AgentSteering::new();
assert_eq!(
steering.try_enqueue("".to_string()),
Err(SteeringRejected::Empty)
);
assert_eq!(
steering.try_enqueue(" \n\t ".to_string()),
Err(SteeringRejected::Empty)
);
assert_eq!(steering.pending_count(), 0);
}
#[test]
fn enqueue_over_capacity_rejects_full() {
let steering = AgentSteering::new();
for index in 0..MAX_PENDING_STEERING_MESSAGES {
assert_eq!(
steering.try_enqueue(format!("message {index}")),
Ok(SteeringQueued {
position: index,
pending_count: index + 1,
})
);
}
assert_eq!(
steering.try_enqueue("overflow".to_string()),
Err(SteeringRejected::Full {
capacity: MAX_PENDING_STEERING_MESSAGES,
})
);
assert_eq!(steering.pending_count(), MAX_PENDING_STEERING_MESSAGES);
}
#[test]
fn acknowledge_observed_prefix_preserves_concurrent_append() {
let steering = AgentSteering::new();
steering.try_enqueue("first".to_string()).unwrap();
steering.try_enqueue("second".to_string()).unwrap();
let batch = steering.observe_collapsed().unwrap();
steering.try_enqueue("later".to_string()).unwrap();
steering.acknowledge(batch.count);
assert_eq!(steering.pending_count(), 1);
assert_eq!(
steering.drain_collapsed(),
Some("Steering update from user while current run was active:\n\nlater".to_string())
);
}
#[test]
fn drain_collapsed_on_empty_queue_returns_none() {
let steering = AgentSteering::new();
assert_eq!(steering.drain_collapsed(), None);
}
#[test]
fn clear_returns_count_and_empties_queue() {
let steering = AgentSteering::new();
steering.try_enqueue("one".to_string()).unwrap();
steering.try_enqueue("two".to_string()).unwrap();
assert_eq!(steering.clear(), 2);
assert_eq!(steering.pending_count(), 0);
assert_eq!(steering.clear(), 0);
}
#[test]
fn recovers_after_queue_lock_poison() {
let steering = AgentSteering::new();
steering.try_enqueue("before poison".to_string()).unwrap();
let _ = std::panic::catch_unwind(|| {
let _guard = steering.inner.lock().unwrap();
panic!("poison steering queue lock");
});
assert_eq!(steering.pending_count(), 1);
assert_eq!(
steering.try_enqueue("after poison".to_string()),
Ok(SteeringQueued {
position: 1,
pending_count: 2,
})
);
assert_eq!(
steering.drain_collapsed(),
Some(
"Steering update from user while current run was active:\n\n1. before poison\n\n2. after poison"
.to_string()
)
);
assert_eq!(steering.clear(), 0);
}
#[test]
fn pending_count_tracks_after_enqueue_and_drain() {
let steering = AgentSteering::new();
assert_eq!(steering.pending_count(), 0);
steering.try_enqueue("one".to_string()).unwrap();
steering.try_enqueue("two".to_string()).unwrap();
assert_eq!(steering.pending_count(), 2);
steering.drain_collapsed().unwrap();
assert_eq!(steering.pending_count(), 0);
}
#[test]
fn drain_collapsed_drains_only_non_empty_queue() {
let steering = AgentSteering::new();
assert_eq!(steering.drain_collapsed(), None);
steering.try_enqueue("one".to_string()).unwrap();
assert_eq!(
steering.drain_collapsed(),
Some("Steering update from user while current run was active:\n\none".to_string())
);
assert_eq!(steering.drain_collapsed(), None);
steering.try_enqueue("two".to_string()).unwrap();
assert_eq!(
steering.drain_collapsed(),
Some("Steering update from user while current run was active:\n\ntwo".to_string())
);
assert_eq!(steering.drain_collapsed(), None);
}
#[test]
fn enqueue_trims_surrounding_whitespace() {
let steering = AgentSteering::new();
steering
.try_enqueue(" trimmed text\n".to_string())
.unwrap();
assert_eq!(
steering.drain_collapsed(),
Some(
"Steering update from user while current run was active:\n\ntrimmed text"
.to_string()
)
);
}
}