use serde::{Deserialize, Serialize}; use serde_json::Value; use std::collections::{HashMap, HashSet, VecDeque}; use std::sync::{Arc, Mutex, MutexGuard}; use thiserror::Error; use uuid::Uuid; /// Delivery mode for a coordination envelope. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] #[serde(rename_all = "snake_case")] pub enum DeliveryScope { /// Send to exactly one target agent. Direct, /// Fan out to all registered agents. Broadcast, } /// Typed payload variants used by agent coordination. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] #[serde(tag = "kind", rename_all = "snake_case")] pub enum CoordinationPayload { DelegateTask { task_id: String, summary: String, metadata: Value, }, ContextPatch { key: String, expected_version: u64, value: Value, }, TaskResult { task_id: String, success: bool, output: String, }, Ack { acked_message_id: String, }, Control { action: String, note: Option, }, } /// Message envelope used by coordination protocol traffic. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct CoordinationEnvelope { pub id: String, pub conversation_id: String, pub correlation_id: Option, pub causation_id: Option, pub from: String, pub to: Option, pub topic: String, pub scope: DeliveryScope, pub payload: CoordinationPayload, } impl CoordinationEnvelope { /// Construct a direct message envelope. pub fn new_direct( from: impl Into, to: impl Into, conversation_id: impl Into, topic: impl Into, payload: CoordinationPayload, ) -> Self { Self { id: Uuid::new_v4().to_string(), conversation_id: conversation_id.into(), correlation_id: None, causation_id: None, from: from.into(), to: Some(to.into()), topic: topic.into(), scope: DeliveryScope::Direct, payload, } } /// Construct a broadcast envelope. pub fn new_broadcast( from: impl Into, conversation_id: impl Into, topic: impl Into, payload: CoordinationPayload, ) -> Self { Self { id: Uuid::new_v4().to_string(), conversation_id: conversation_id.into(), correlation_id: None, causation_id: None, from: from.into(), to: None, topic: topic.into(), scope: DeliveryScope::Broadcast, payload, } } /// Validate transport and payload contract before publishing. pub fn validate(&self) -> Result<(), CoordinationError> { require_non_empty(&self.id, "id")?; require_non_empty(&self.conversation_id, "conversation_id")?; require_non_empty(&self.from, "from")?; require_non_empty(&self.topic, "topic")?; match self.scope { DeliveryScope::Direct => { let target = self .to .as_deref() .map(str::trim) .filter(|value| !value.is_empty()); if target.is_none() { return Err(CoordinationError::MissingTarget { message_id: self.id.clone(), }); } } DeliveryScope::Broadcast => { if self.to.is_some() { return Err(CoordinationError::BroadcastHasTarget { message_id: self.id.clone(), }); } } } if let Some(correlation_id) = &self.correlation_id { require_non_empty(correlation_id, "correlation_id")?; } if let Some(causation_id) = &self.causation_id { require_non_empty(causation_id, "causation_id")?; } match &self.payload { CoordinationPayload::DelegateTask { task_id, summary, .. } => { require_non_empty(task_id, "task_id")?; require_non_empty(summary, "summary")?; if self.scope != DeliveryScope::Direct { return Err(CoordinationError::InvalidDeliveryScope { message_id: self.id.clone(), expected: DeliveryScope::Direct, actual: self.scope, payload: "delegate_task".to_string(), }); } } CoordinationPayload::ContextPatch { key, .. } => { require_non_empty(key, "key")?; } CoordinationPayload::TaskResult { task_id, output, .. } => { require_non_empty(task_id, "task_id")?; require_non_empty(output, "output")?; if self .correlation_id .as_deref() .map(str::trim) .filter(|value| !value.is_empty()) .is_none() { return Err(CoordinationError::MissingCorrelationId { message_id: self.id.clone(), }); } } CoordinationPayload::Ack { acked_message_id } => { require_non_empty(acked_message_id, "acked_message_id")?; } CoordinationPayload::Control { action, .. } => { require_non_empty(action, "action")?; } } Ok(()) } } fn require_non_empty(value: &str, field: &'static str) -> Result<(), CoordinationError> { if value.trim().is_empty() { return Err(CoordinationError::EmptyField { field }); } Ok(()) } /// Errors emitted by the coordination protocol and message bus. #[derive(Debug, Error, Clone, PartialEq, Eq)] pub enum CoordinationError { #[error("field `{field}` must not be empty")] EmptyField { field: &'static str }, #[error("message `{message_id}` requires a direct target agent")] MissingTarget { message_id: String }, #[error("broadcast message `{message_id}` cannot set explicit target")] BroadcastHasTarget { message_id: String }, #[error("task result message `{message_id}` requires `correlation_id`")] MissingCorrelationId { message_id: String }, #[error( "invalid delivery scope for payload `{payload}` on message `{message_id}`: expected {expected:?}, got {actual:?}" )] InvalidDeliveryScope { message_id: String, expected: DeliveryScope, actual: DeliveryScope, payload: String, }, #[error("duplicate message id `{message_id}`")] DuplicateMessageId { message_id: String }, #[error("unknown target agent `{agent}` for message `{message_id}`")] UnknownTarget { agent: String, message_id: String }, #[error("agent `{agent}` is not registered")] UnknownAgent { agent: String }, #[error("invalid delegate context key `{key}` on message `{message_id}`")] InvalidDelegateContextKey { key: String, message_id: String }, #[error("delegate context key `{key}` requires `correlation_id` on message `{message_id}`")] MissingDelegateContextCorrelation { key: String, message_id: String }, #[error( "delegate context key `{key}` correlation mismatch on message `{message_id}`: key has `{key_correlation_id}`, envelope has `{envelope_correlation_id}`" )] DelegateContextCorrelationMismatch { key: String, message_id: String, key_correlation_id: String, envelope_correlation_id: String, }, #[error("context version mismatch for key `{key}`: expected {expected}, actual {actual}")] ContextVersionMismatch { key: String, expected: u64, actual: u64, }, } /// Sequenced message emitted by the bus. #[derive(Debug, Clone)] pub struct SequencedEnvelope { pub sequence: u64, pub envelope: CoordinationEnvelope, } /// Dead-letter item retained for audit and debugging. #[derive(Debug, Clone)] pub struct DeadLetter { pub envelope: CoordinationEnvelope, pub reason: String, } /// Versioned shared context record written through `ContextPatch`. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct SharedContextEntry { pub key: String, pub value: Value, pub version: u64, pub updated_by: String, pub last_message_id: String, } /// Publish result metadata. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct PublishReceipt { pub sequence: u64, pub delivered_to: usize, } /// Capacity limits used by `InMemoryMessageBus` retention policies. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] pub struct InMemoryMessageBusLimits { pub max_inbox_messages_per_agent: usize, pub max_dead_letters: usize, pub max_context_entries: usize, pub max_seen_message_ids: usize, } impl Default for InMemoryMessageBusLimits { fn default() -> Self { Self { max_inbox_messages_per_agent: 256, max_dead_letters: 256, max_context_entries: 512, max_seen_message_ids: 4096, } } } /// Runtime counters for operational visibility. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)] pub struct InMemoryMessageBusStats { /// Total publish attempts that passed envelope validation. pub publish_attempts_total: u64, /// Total successful deliveries (fan-out count for broadcast). pub deliveries_total: u64, /// Number of inbox messages evicted due to inbox capacity limits. pub inbox_overflow_evictions_total: u64, /// Total dead-letter entries ever recorded. pub dead_letters_total: u64, /// Number of dead-letter entries evicted due to dead-letter cap. pub dead_letter_evictions_total: u64, /// Number of shared-context entries evicted due to context capacity limits. pub context_evictions_total: u64, /// Number of idempotency IDs evicted due to dedupe-window capacity limits. pub seen_message_id_evictions_total: u64, } #[derive(Debug, Default)] struct BusState { next_sequence: u64, seen_message_ids: HashSet, seen_message_order: VecDeque, inboxes: HashMap>, inbox_correlation_counts: HashMap>, dead_letters: Vec, dead_letters_by_correlation: HashMap>, context: HashMap, context_order: VecDeque, delegate_context_order: VecDeque, context_order_by_correlation: HashMap>, delegate_context_order_by_correlation: HashMap>, context_correlation_by_key: HashMap, limits: InMemoryMessageBusLimits, stats: InMemoryMessageBusStats, } impl BusState { fn with_limits(mut limits: InMemoryMessageBusLimits) -> Self { if limits.max_inbox_messages_per_agent == 0 { limits.max_inbox_messages_per_agent = 1; } if limits.max_dead_letters == 0 { limits.max_dead_letters = 1; } if limits.max_context_entries == 0 { limits.max_context_entries = 1; } if limits.max_seen_message_ids == 0 { limits.max_seen_message_ids = 1; } Self { next_sequence: 0, seen_message_ids: HashSet::new(), seen_message_order: VecDeque::new(), inboxes: HashMap::new(), inbox_correlation_counts: HashMap::new(), dead_letters: Vec::new(), dead_letters_by_correlation: HashMap::new(), context: HashMap::new(), context_order: VecDeque::new(), delegate_context_order: VecDeque::new(), context_order_by_correlation: HashMap::new(), delegate_context_order_by_correlation: HashMap::new(), context_correlation_by_key: HashMap::new(), limits, stats: InMemoryMessageBusStats::default(), } } } /// Deterministic in-memory coordination message bus with: /// - typed envelope validation /// - idempotency guard on message id /// - per-agent ordered delivery /// - dead-letter retention for invalid/conflicting messages /// - optimistic-locking context patches #[derive(Debug, Clone)] pub struct InMemoryMessageBus { inner: Arc>, } impl Default for InMemoryMessageBus { fn default() -> Self { Self::new() } } impl InMemoryMessageBus { pub fn new() -> Self { Self::with_limits(InMemoryMessageBusLimits::default()) } pub fn with_limits(limits: InMemoryMessageBusLimits) -> Self { Self { inner: Arc::new(Mutex::new(BusState::with_limits(limits))), } } /// Register an agent inbox. pub fn register_agent(&self, agent: impl Into) -> Result<(), CoordinationError> { let agent = agent.into(); require_non_empty(&agent, "agent")?; let mut state = self.lock_state(); state.inboxes.entry(agent.clone()).or_default(); state.inbox_correlation_counts.entry(agent).or_default(); Ok(()) } /// Remove an existing agent inbox. pub fn unregister_agent(&self, agent: &str) -> bool { let mut state = self.lock_state(); let removed = state.inboxes.remove(agent).is_some(); state.inbox_correlation_counts.remove(agent); removed } /// Publish an envelope to the bus. pub fn publish( &self, envelope: CoordinationEnvelope, ) -> Result { if let Err(error) = envelope.validate() { self.push_dead_letter(envelope, error.to_string()); return Err(error); } let mut state = self.lock_state(); state.stats.publish_attempts_total += 1; if state.seen_message_ids.contains(&envelope.id) { let error = CoordinationError::DuplicateMessageId { message_id: envelope.id.clone(), }; push_dead_letter_locked(&mut state, envelope, error.to_string()); return Err(error); } if state.seen_message_ids.len() >= state.limits.max_seen_message_ids { if let Some(evicted_id) = state.seen_message_order.pop_front() { if state.seen_message_ids.remove(&evicted_id) { state.stats.seen_message_id_evictions_total += 1; } } } state.seen_message_ids.insert(envelope.id.clone()); state.seen_message_order.push_back(envelope.id.clone()); if let CoordinationPayload::ContextPatch { key, expected_version, value, } = &envelope.payload { if let Err(error) = apply_context_patch_locked(&mut state, &envelope, key, *expected_version, value) { push_dead_letter_locked(&mut state, envelope, error.to_string()); return Err(error); } } state.next_sequence += 1; let sequence = state.next_sequence; let sequenced = SequencedEnvelope { sequence, envelope: envelope.clone(), }; let delivered_to = match envelope.scope { DeliveryScope::Direct => { let target = envelope.to.as_deref().expect("validated direct target"); if !state.inboxes.contains_key(target) { let error = CoordinationError::UnknownTarget { agent: target.to_string(), message_id: envelope.id.clone(), }; push_dead_letter_locked(&mut state, envelope, error.to_string()); return Err(error); } let dropped = push_inbox_entry_locked(&mut state, target, sequenced); if let Some(dropped) = dropped { state.stats.inbox_overflow_evictions_total += 1; push_dead_letter_locked( &mut state, dropped, format!("inbox overflow: dropped oldest message for agent '{target}'"), ); } 1 } DeliveryScope::Broadcast => { if state.inboxes.is_empty() { 0 } else { let fanout = state.inboxes.len(); let mut dropped_items: Vec<(String, CoordinationEnvelope)> = Vec::new(); let agents = state.inboxes.keys().cloned().collect::>(); for agent in &agents { if let Some(dropped) = push_inbox_entry_locked(&mut state, agent, sequenced.clone()) { dropped_items.push((agent.clone(), dropped)); } } for (agent, dropped) in dropped_items { state.stats.inbox_overflow_evictions_total += 1; push_dead_letter_locked( &mut state, dropped, format!("inbox overflow: dropped oldest message for agent '{agent}'"), ); } fanout } } }; state.stats.deliveries_total += delivered_to as u64; Ok(PublishReceipt { sequence, delivered_to, }) } /// Drain up to `max` pending envelopes for an agent inbox. /// Use `max = 0` to drain all available messages. pub fn drain_for_agent( &self, agent: &str, max: usize, ) -> Result, CoordinationError> { let mut state = self.lock_state(); let agent_owned = agent.to_string(); let inbox_len = state.inboxes.get(agent).map(VecDeque::len).ok_or_else(|| { CoordinationError::UnknownAgent { agent: agent_owned.clone(), } })?; let drain_count = if max == 0 { inbox_len } else { max.min(inbox_len) }; let mut drained = Vec::with_capacity(drain_count); for _ in 0..drain_count { let envelope = { let inbox = state .inboxes .get_mut(agent) .expect("agent existence should be validated before drain"); inbox.pop_front() }; if let Some(envelope) = envelope { let correlation_counts = state .inbox_correlation_counts .entry(agent_owned.clone()) .or_default(); decrement_correlation_count(correlation_counts, &envelope.envelope); drained.push(envelope); } } Ok(drained) } pub fn pending_for_agent(&self, agent: &str) -> Result { let state = self.lock_state(); state .inboxes .get(agent) .map(VecDeque::len) .ok_or_else(|| CoordinationError::UnknownAgent { agent: agent.to_string(), }) } pub fn pending_for_agent_correlation( &self, agent: &str, correlation_id: &str, ) -> Result { let correlation_id = correlation_id.trim(); if correlation_id.is_empty() { return Ok(0); } let state = self.lock_state(); if !state.inboxes.contains_key(agent) { return Err(CoordinationError::UnknownAgent { agent: agent.to_string(), }); } Ok(state .inbox_correlation_counts .get(agent) .and_then(|counts| counts.get(correlation_id).copied()) .unwrap_or(0)) } /// Peek up to `max` pending envelopes for an agent without consuming them. /// Use `max = 0` to peek the full inbox. pub fn peek_for_agent( &self, agent: &str, max: usize, ) -> Result, CoordinationError> { self.peek_for_agent_with_offset(agent, 0, max) } /// Peek up to `max` pending envelopes for an agent without consuming them, /// with an offset into inbox order (oldest first). /// Use `max = 0` to peek all entries after `offset`. pub fn peek_for_agent_with_offset( &self, agent: &str, offset: usize, max: usize, ) -> Result, CoordinationError> { let state = self.lock_state(); let inbox = state .inboxes .get(agent) .ok_or_else(|| CoordinationError::UnknownAgent { agent: agent.to_string(), })?; let available = inbox.len().saturating_sub(offset); let take_count = if max == 0 { available } else { max.min(available) }; Ok(inbox .iter() .skip(offset) .take(take_count) .cloned() .collect()) } /// Peek up to `max` pending envelopes matching a correlation ID for an /// agent without consuming them, with an offset in match order /// (oldest first). Use `max = 0` to peek all matches after `offset`. pub fn peek_for_agent_correlation_with_offset( &self, agent: &str, correlation_id: &str, offset: usize, max: usize, ) -> Result, CoordinationError> { let correlation_id = correlation_id.trim(); if correlation_id.is_empty() { return Ok(Vec::new()); } let state = self.lock_state(); let inbox = state .inboxes .get(agent) .ok_or_else(|| CoordinationError::UnknownAgent { agent: agent.to_string(), })?; let available = state .inbox_correlation_counts .get(agent) .and_then(|counts| counts.get(correlation_id).copied()) .unwrap_or(0) .saturating_sub(offset); let take_count = if max == 0 { available } else { max.min(available) }; Ok(inbox .iter() .filter(|entry| { normalized_non_empty(entry.envelope.correlation_id.as_deref()) .is_some_and(|value| value == correlation_id) }) .skip(offset) .take(take_count) .cloned() .collect()) } /// Snapshot registered agents with inboxes. pub fn registered_agents(&self) -> Vec { let state = self.lock_state(); let mut agents = state.inboxes.keys().cloned().collect::>(); agents.sort(); agents } pub fn limits(&self) -> InMemoryMessageBusLimits { self.lock_state().limits } pub fn stats(&self) -> InMemoryMessageBusStats { self.lock_state().stats } pub fn subscriber_count(&self) -> usize { self.lock_state().inboxes.len() } /// Snapshot all shared context entries. pub fn context_snapshot(&self) -> HashMap { self.lock_state().context.clone() } /// Snapshot shared context entries in write-recency order (newest first). /// Use `max = 0` to return all entries. pub fn context_entries_recent(&self, max: usize) -> Vec<(String, SharedContextEntry)> { self.context_entries_recent_with_offset(0, max) } /// Snapshot shared context entries in write-recency order (newest first), /// with an offset for pagination. /// Use `max = 0` to return all entries after `offset`. pub fn context_entries_recent_with_offset( &self, offset: usize, max: usize, ) -> Vec<(String, SharedContextEntry)> { let state = self.lock_state(); let available = state.context_order.len().saturating_sub(offset); let take_count = if max == 0 { available } else { max.min(available) }; state .context_order .iter() .rev() .skip(offset) .take(take_count) .filter_map(|key| { state .context .get(key) .cloned() .map(|entry| (key.clone(), entry)) }) .collect() } /// Snapshot shared context entries for a correlation ID in write-recency /// order (newest first). Use `max = 0` to return all entries. pub fn context_entries_recent_for_correlation( &self, correlation_id: &str, max: usize, ) -> Vec<(String, SharedContextEntry)> { self.context_entries_recent_for_correlation_with_offset(correlation_id, 0, max) } /// Snapshot shared context entries for a correlation ID in write-recency /// order (newest first), with an offset for pagination. /// Use `max = 0` to return all entries after `offset`. pub fn context_entries_recent_for_correlation_with_offset( &self, correlation_id: &str, offset: usize, max: usize, ) -> Vec<(String, SharedContextEntry)> { let correlation_id = correlation_id.trim(); if correlation_id.is_empty() { return Vec::new(); } let state = self.lock_state(); let Some(order) = state.context_order_by_correlation.get(correlation_id) else { return Vec::new(); }; let available = order.len().saturating_sub(offset); let take_count = if max == 0 { available } else { max.min(available) }; order .iter() .rev() .skip(offset) .take(take_count) .filter_map(|key| { state .context .get(key) .cloned() .map(|entry| (key.clone(), entry)) }) .collect() } pub fn context_count(&self) -> usize { self.lock_state().context.len() } pub fn context_count_for_correlation(&self, correlation_id: &str) -> usize { let correlation_id = correlation_id.trim(); if correlation_id.is_empty() { return 0; } let state = self.lock_state(); state .context_order_by_correlation .get(correlation_id) .map(VecDeque::len) .unwrap_or(0) } /// Snapshot only `delegate/` shared context entries in write-recency order /// (newest first), with an offset for pagination. /// Use `max = 0` to return all entries after `offset`. pub fn delegate_context_entries_recent_with_offset( &self, offset: usize, max: usize, ) -> Vec<(String, SharedContextEntry)> { let state = self.lock_state(); let available = state.delegate_context_order.len().saturating_sub(offset); let take_count = if max == 0 { available } else { max.min(available) }; state .delegate_context_order .iter() .rev() .skip(offset) .take(take_count) .filter_map(|key| { state .context .get(key) .cloned() .map(|entry| (key.clone(), entry)) }) .collect() } /// Snapshot only `delegate/` shared context entries for a correlation ID /// in write-recency order (newest first), with an offset for pagination. /// Use `max = 0` to return all entries after `offset`. pub fn delegate_context_entries_recent_for_correlation_with_offset( &self, correlation_id: &str, offset: usize, max: usize, ) -> Vec<(String, SharedContextEntry)> { let correlation_id = correlation_id.trim(); if correlation_id.is_empty() { return Vec::new(); } let state = self.lock_state(); let Some(order) = state .delegate_context_order_by_correlation .get(correlation_id) else { return Vec::new(); }; let available = order.len().saturating_sub(offset); let take_count = if max == 0 { available } else { max.min(available) }; order .iter() .rev() .skip(offset) .take(take_count) .filter_map(|key| { state .context .get(key) .cloned() .map(|entry| (key.clone(), entry)) }) .collect() } pub fn delegate_context_count(&self) -> usize { self.lock_state().delegate_context_order.len() } pub fn delegate_context_count_for_correlation(&self, correlation_id: &str) -> usize { let correlation_id = correlation_id.trim(); if correlation_id.is_empty() { return 0; } let state = self.lock_state(); state .delegate_context_order_by_correlation .get(correlation_id) .map(VecDeque::len) .unwrap_or(0) } /// Snapshot dead-letter entries in recency order (newest first), /// with an offset for pagination. /// Use `max = 0` to return all entries after `offset`. pub fn dead_letters_recent(&self, offset: usize, max: usize) -> Vec { let state = self.lock_state(); let available = state.dead_letters.len().saturating_sub(offset); let take_count = if max == 0 { available } else { max.min(available) }; state .dead_letters .iter() .rev() .skip(offset) .take(take_count) .cloned() .collect() } /// Snapshot dead-letter entries for a correlation ID in recency order /// (newest first), with an offset for pagination. /// Use `max = 0` to return all entries after `offset`. pub fn dead_letters_recent_for_correlation( &self, correlation_id: &str, offset: usize, max: usize, ) -> Vec { let correlation_id = correlation_id.trim(); if correlation_id.is_empty() { return Vec::new(); } let state = self.lock_state(); let Some(entries) = state.dead_letters_by_correlation.get(correlation_id) else { return Vec::new(); }; let available = entries.len().saturating_sub(offset); let take_count = if max == 0 { available } else { max.min(available) }; entries .iter() .rev() .skip(offset) .take(take_count) .cloned() .collect() } pub fn context_entry(&self, key: &str) -> Option { self.lock_state().context.get(key).cloned() } pub fn dead_letter_count(&self) -> usize { self.lock_state().dead_letters.len() } pub fn dead_letter_count_for_correlation(&self, correlation_id: &str) -> usize { let correlation_id = correlation_id.trim(); if correlation_id.is_empty() { return 0; } let state = self.lock_state(); state .dead_letters_by_correlation .get(correlation_id) .map(VecDeque::len) .unwrap_or(0) } pub fn dead_letters(&self) -> Vec { self.lock_state().dead_letters.clone() } fn push_dead_letter(&self, envelope: CoordinationEnvelope, reason: String) { let mut state = self.lock_state(); push_dead_letter_locked(&mut state, envelope, reason); } fn lock_state(&self) -> MutexGuard<'_, BusState> { self.inner .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) } } fn push_inbox_entry_locked( state: &mut BusState, agent: &str, entry: SequencedEnvelope, ) -> Option { let max_inbox_messages_per_agent = state.limits.max_inbox_messages_per_agent; let (inboxes, correlation_counts_by_agent) = (&mut state.inboxes, &mut state.inbox_correlation_counts); let inbox = inboxes .get_mut(agent) .expect("agent existence should be validated before pushing inbox entry"); let correlation_counts = correlation_counts_by_agent .entry(agent.to_string()) .or_default(); let dropped = if inbox.len() >= max_inbox_messages_per_agent { inbox.pop_front() } else { None }; if let Some(dropped_entry) = dropped.as_ref() { decrement_correlation_count(correlation_counts, &dropped_entry.envelope); } increment_correlation_count(correlation_counts, &entry.envelope); inbox.push_back(entry); dropped.map(|value| value.envelope) } fn increment_correlation_count( counts: &mut HashMap, envelope: &CoordinationEnvelope, ) { if let Some(correlation_id) = normalized_non_empty(envelope.correlation_id.as_deref()) { *counts.entry(correlation_id.to_string()).or_insert(0) += 1; } } fn decrement_correlation_count( counts: &mut HashMap, envelope: &CoordinationEnvelope, ) { let Some(correlation_id) = normalized_non_empty(envelope.correlation_id.as_deref()) else { return; }; let mut remove_key = false; if let Some(count) = counts.get_mut(correlation_id) { if *count <= 1 { remove_key = true; } else { *count -= 1; } } if remove_key { counts.remove(correlation_id); } } fn push_dead_letter_locked(state: &mut BusState, envelope: CoordinationEnvelope, reason: String) { state.stats.dead_letters_total += 1; if state.dead_letters.len() >= state.limits.max_dead_letters { state.stats.dead_letter_evictions_total += 1; if let Some(evicted) = state.dead_letters.first() { if let Some(correlation_id) = normalized_non_empty(evicted.envelope.correlation_id.as_deref()) { let mut remove_correlation_key = false; if let Some(entries) = state.dead_letters_by_correlation.get_mut(correlation_id) { let _ = entries.pop_front(); remove_correlation_key = entries.is_empty(); } if remove_correlation_key { state.dead_letters_by_correlation.remove(correlation_id); } } } let _ = state.dead_letters.remove(0); } let dead_letter = DeadLetter { envelope, reason }; if let Some(correlation_id) = normalized_non_empty(dead_letter.envelope.correlation_id.as_deref()) { state .dead_letters_by_correlation .entry(correlation_id.to_string()) .or_default() .push_back(dead_letter.clone()); } state.dead_letters.push(dead_letter); } fn apply_context_patch_locked( state: &mut BusState, envelope: &CoordinationEnvelope, key: &str, expected_version: u64, value: &Value, ) -> Result<(), CoordinationError> { let key_delegate_correlation = if key.starts_with("delegate/") { let parsed = parse_delegate_context_correlation_from_key(key).ok_or_else(|| { CoordinationError::InvalidDelegateContextKey { key: key.to_string(), message_id: envelope.id.clone(), } })?; let envelope_correlation = normalized_non_empty(envelope.correlation_id.as_deref()) .ok_or_else(|| CoordinationError::MissingDelegateContextCorrelation { key: key.to_string(), message_id: envelope.id.clone(), })?; if parsed != envelope_correlation { return Err(CoordinationError::DelegateContextCorrelationMismatch { key: key.to_string(), message_id: envelope.id.clone(), key_correlation_id: parsed.to_string(), envelope_correlation_id: envelope_correlation.to_string(), }); } Some(parsed) } else { None }; let current_version = state.context.get(key).map_or(0, |entry| entry.version); if current_version != expected_version { return Err(CoordinationError::ContextVersionMismatch { key: key.to_string(), expected: expected_version, actual: current_version, }); } let key_owned = key.to_string(); let key_is_delegate = key_delegate_correlation.is_some(); let previous_correlation = state.context_correlation_by_key.get(key).cloned(); let is_new_key = !state.context.contains_key(key); if is_new_key && state.context.len() >= state.limits.max_context_entries { if let Some(evicted_key) = state.context_order.pop_front() { if state.context.remove(&evicted_key).is_some() { state.stats.context_evictions_total += 1; } let evicted_correlation = state.context_correlation_by_key.remove(&evicted_key); if let Some(correlation_id) = evicted_correlation.as_deref() { remove_key_from_context_correlation_order(state, correlation_id, &evicted_key); } if evicted_key.starts_with("delegate/") { remove_key_from_delegate_context_order( state, &evicted_key, evicted_correlation.as_deref(), ); } } } if !is_new_key { if let Some(position) = state .context_order .iter() .position(|existing| existing == key) { let _ = state.context_order.remove(position); } } state.context_order.push_back(key_owned.clone()); if let Some(correlation_id) = previous_correlation.as_deref() { remove_key_from_context_correlation_order(state, correlation_id, key); } if key_is_delegate { remove_key_from_delegate_context_order(state, key, previous_correlation.as_deref()); state.delegate_context_order.push_back(key_owned.clone()); } if let Some(correlation_id) = normalized_non_empty(envelope.correlation_id.as_deref()) { state .context_order_by_correlation .entry(correlation_id.to_string()) .or_default() .push_back(key_owned.clone()); if key_is_delegate { state .delegate_context_order_by_correlation .entry(correlation_id.to_string()) .or_default() .push_back(key_owned.clone()); } state .context_correlation_by_key .insert(key_owned.clone(), correlation_id.to_string()); } else { state.context_correlation_by_key.remove(&key_owned); } state.context.insert( key_owned.clone(), SharedContextEntry { key: key_owned, value: value.clone(), version: current_version + 1, updated_by: envelope.from.clone(), last_message_id: envelope.id.clone(), }, ); Ok(()) } fn remove_key_from_context_correlation_order( state: &mut BusState, correlation_id: &str, key: &str, ) { let mut remove_correlation_key = false; if let Some(order) = state.context_order_by_correlation.get_mut(correlation_id) { if let Some(position) = order.iter().position(|existing| existing == key) { let _ = order.remove(position); } remove_correlation_key = order.is_empty(); } if remove_correlation_key { state.context_order_by_correlation.remove(correlation_id); } } fn remove_key_from_delegate_context_order( state: &mut BusState, key: &str, correlation_id: Option<&str>, ) { if let Some(position) = state .delegate_context_order .iter() .position(|existing| existing == key) { let _ = state.delegate_context_order.remove(position); } let Some(correlation_id) = correlation_id else { return; }; let mut remove_correlation_key = false; if let Some(order) = state .delegate_context_order_by_correlation .get_mut(correlation_id) { if let Some(position) = order.iter().position(|existing| existing == key) { let _ = order.remove(position); } remove_correlation_key = order.is_empty(); } if remove_correlation_key { state .delegate_context_order_by_correlation .remove(correlation_id); } } fn normalized_non_empty(value: Option<&str>) -> Option<&str> { value.map(str::trim).filter(|value| !value.is_empty()) } fn parse_delegate_context_correlation_from_key(key: &str) -> Option<&str> { let mut parts = key.splitn(3, '/'); let namespace = parts.next()?; if namespace != "delegate" { return None; } let correlation = parts.next()?.trim(); if correlation.is_empty() { return None; } // Require at least one trailing segment (e.g. delegate//state). let tail = parts.next()?.trim(); if tail.is_empty() { return None; } Some(correlation) } #[cfg(test)] mod tests { use super::*; use serde_json::json; use std::collections::HashSet; use tokio::sync::Barrier; #[test] fn delegate_task_requires_direct_target() { let envelope = CoordinationEnvelope { id: "msg-1".to_string(), conversation_id: "conv-1".to_string(), correlation_id: None, causation_id: None, from: "lead".to_string(), to: None, topic: "coordination".to_string(), scope: DeliveryScope::Direct, payload: CoordinationPayload::DelegateTask { task_id: "task-1".to_string(), summary: "Investigate bug".to_string(), metadata: json!({}), }, }; let error = envelope .validate() .expect_err("target agent must be required"); assert_eq!( error, CoordinationError::MissingTarget { message_id: "msg-1".to_string() } ); } #[test] fn task_result_requires_correlation_id() { let envelope = CoordinationEnvelope { id: "msg-2".to_string(), conversation_id: "conv-1".to_string(), correlation_id: None, causation_id: None, from: "worker".to_string(), to: Some("lead".to_string()), topic: "coordination".to_string(), scope: DeliveryScope::Direct, payload: CoordinationPayload::TaskResult { task_id: "task-1".to_string(), success: true, output: "done".to_string(), }, }; let error = envelope .validate() .expect_err("task result must require correlation"); assert_eq!( error, CoordinationError::MissingCorrelationId { message_id: "msg-2".to_string() } ); } #[test] fn json_roundtrip_keeps_payload_shape() { let mut envelope = CoordinationEnvelope::new_direct( "lead", "worker", "conv-1", "coordination", CoordinationPayload::DelegateTask { task_id: "task-1".to_string(), summary: "Analyze logs".to_string(), metadata: json!({"priority": "high"}), }, ); envelope.correlation_id = Some("corr-1".to_string()); let encoded = serde_json::to_string(&envelope).expect("serialize envelope"); let decoded: CoordinationEnvelope = serde_json::from_str(&encoded).expect("deserialize envelope"); assert_eq!(decoded, envelope); } #[test] fn duplicate_message_ids_are_rejected_and_dead_lettered() { let bus = InMemoryMessageBus::new(); bus.register_agent("worker").expect("register worker"); let mut envelope = CoordinationEnvelope::new_direct( "lead", "worker", "conv-1", "coordination", CoordinationPayload::DelegateTask { task_id: "task-1".to_string(), summary: "Investigate".to_string(), metadata: json!({}), }, ); envelope.id = "fixed-id".to_string(); let first = bus.publish(envelope.clone()).expect("first publish"); assert_eq!(first.delivered_to, 1); let second = bus.publish(envelope).expect_err("duplicate id must fail"); assert_eq!( second, CoordinationError::DuplicateMessageId { message_id: "fixed-id".to_string() } ); let dead_letters = bus.dead_letters(); assert_eq!(dead_letters.len(), 1); assert!(dead_letters[0].reason.contains("duplicate message id")); let stats = bus.stats(); assert_eq!(stats.seen_message_id_evictions_total, 0); } #[test] fn dedupe_window_evicts_old_ids_and_allows_reuse_after_eviction() { let bus = InMemoryMessageBus::with_limits(InMemoryMessageBusLimits { max_inbox_messages_per_agent: 32, max_dead_letters: 32, max_context_entries: 32, max_seen_message_ids: 2, }); bus.register_agent("worker").expect("register worker"); for message_id in ["msg-0", "msg-1", "msg-2"] { let mut envelope = CoordinationEnvelope::new_direct( "lead", "worker", "conv-dedupe-window", "coordination", CoordinationPayload::DelegateTask { task_id: message_id.to_string(), summary: "Investigate".to_string(), metadata: json!({}), }, ); envelope.id = message_id.to_string(); bus.publish(envelope).expect("publish should succeed"); } // `msg-0` has been evicted from dedupe window and can be reused. let mut reused = CoordinationEnvelope::new_direct( "lead", "worker", "conv-dedupe-window", "coordination", CoordinationPayload::DelegateTask { task_id: "msg-0".to_string(), summary: "Investigate again".to_string(), metadata: json!({}), }, ); reused.id = "msg-0".to_string(); bus.publish(reused) .expect("reused id should be accepted after eviction"); // Recent IDs are still protected by dedupe window. let mut duplicate_recent = CoordinationEnvelope::new_direct( "lead", "worker", "conv-dedupe-window", "coordination", CoordinationPayload::DelegateTask { task_id: "msg-2".to_string(), summary: "duplicate".to_string(), metadata: json!({}), }, ); duplicate_recent.id = "msg-2".to_string(); let error = bus .publish(duplicate_recent) .expect_err("recent duplicate should be rejected"); assert_eq!( error, CoordinationError::DuplicateMessageId { message_id: "msg-2".to_string() } ); let stats = bus.stats(); assert_eq!(stats.seen_message_id_evictions_total, 2); } #[test] fn context_patch_conflict_goes_to_dead_letter() { let bus = InMemoryMessageBus::new(); bus.register_agent("lead").expect("register lead"); let first_patch = CoordinationEnvelope::new_broadcast( "lead", "conv-ctx", "context", CoordinationPayload::ContextPatch { key: "task-99/state".to_string(), expected_version: 0, value: json!({"phase": "started"}), }, ); bus.publish(first_patch).expect("first patch must succeed"); let stale_patch = CoordinationEnvelope::new_broadcast( "lead", "conv-ctx", "context", CoordinationPayload::ContextPatch { key: "task-99/state".to_string(), expected_version: 0, value: json!({"phase": "stale"}), }, ); let error = bus .publish(stale_patch) .expect_err("stale expected_version must fail"); assert_eq!( error, CoordinationError::ContextVersionMismatch { key: "task-99/state".to_string(), expected: 0, actual: 1 } ); let entry = bus .context_entry("task-99/state") .expect("context entry must exist"); assert_eq!(entry.version, 1); assert_eq!(entry.value, json!({"phase": "started"})); assert_eq!(bus.dead_letters().len(), 1); } #[test] fn delegate_context_patch_requires_correlation_id() { let bus = InMemoryMessageBus::new(); let mut patch = CoordinationEnvelope::new_broadcast( "lead", "conv-delegate-context-correlation", "delegate.state", CoordinationPayload::ContextPatch { key: "delegate/corr-a/state".to_string(), expected_version: 0, value: json!({"phase":"queued"}), }, ); patch.id = "msg-delegate-corr-required".to_string(); let error = bus .publish(patch) .expect_err("delegate context patch without correlation must fail"); assert_eq!( error, CoordinationError::MissingDelegateContextCorrelation { key: "delegate/corr-a/state".to_string(), message_id: "msg-delegate-corr-required".to_string(), } ); assert_eq!(bus.dead_letter_count(), 1); } #[test] fn delegate_context_patch_rejects_mismatched_correlation_id() { let bus = InMemoryMessageBus::new(); let mut patch = CoordinationEnvelope::new_broadcast( "lead", "conv-delegate-context-correlation-mismatch", "delegate.state", CoordinationPayload::ContextPatch { key: "delegate/corr-a/state".to_string(), expected_version: 0, value: json!({"phase":"queued"}), }, ); patch.id = "msg-delegate-corr-mismatch".to_string(); patch.correlation_id = Some("corr-b".to_string()); let error = bus .publish(patch) .expect_err("delegate context patch with mismatch must fail"); assert_eq!( error, CoordinationError::DelegateContextCorrelationMismatch { key: "delegate/corr-a/state".to_string(), message_id: "msg-delegate-corr-mismatch".to_string(), key_correlation_id: "corr-a".to_string(), envelope_correlation_id: "corr-b".to_string(), } ); assert_eq!(bus.dead_letter_count(), 1); } #[test] fn delegate_context_patch_rejects_invalid_delegate_key_shape() { let bus = InMemoryMessageBus::new(); let mut patch = CoordinationEnvelope::new_broadcast( "lead", "conv-delegate-context-key-shape", "delegate.state", CoordinationPayload::ContextPatch { key: "delegate/corr-a".to_string(), expected_version: 0, value: json!({"phase":"queued"}), }, ); patch.id = "msg-delegate-key-shape".to_string(); patch.correlation_id = Some("corr-a".to_string()); let error = bus .publish(patch) .expect_err("delegate context patch with invalid key shape must fail"); assert_eq!( error, CoordinationError::InvalidDelegateContextKey { key: "delegate/corr-a".to_string(), message_id: "msg-delegate-key-shape".to_string(), } ); assert_eq!(bus.dead_letter_count(), 1); } #[test] fn delegate_context_patch_rejects_empty_tail_segment() { let bus = InMemoryMessageBus::new(); let mut patch = CoordinationEnvelope::new_broadcast( "lead", "conv-delegate-context-key-tail", "delegate.state", CoordinationPayload::ContextPatch { key: "delegate/corr-a/".to_string(), expected_version: 0, value: json!({"phase":"queued"}), }, ); patch.id = "msg-delegate-key-tail".to_string(); patch.correlation_id = Some("corr-a".to_string()); let error = bus .publish(patch) .expect_err("delegate context patch with empty tail must fail"); assert_eq!( error, CoordinationError::InvalidDelegateContextKey { key: "delegate/corr-a/".to_string(), message_id: "msg-delegate-key-tail".to_string(), } ); assert_eq!(bus.dead_letter_count(), 1); } #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn concurrent_publish_keeps_inbox_order() { let bus = InMemoryMessageBus::new(); bus.register_agent("lead").expect("register lead"); bus.register_agent("worker").expect("register worker"); let total = 32usize; let barrier = Arc::new(Barrier::new(total)); let mut tasks = Vec::with_capacity(total); for index in 0..total { let bus_clone = bus.clone(); let barrier_clone = Arc::clone(&barrier); tasks.push(tokio::spawn(async move { barrier_clone.wait().await; let mut envelope = CoordinationEnvelope::new_direct( "lead", "worker", "conv-concurrent", "coordination", CoordinationPayload::DelegateTask { task_id: format!("task-{index}"), summary: format!("work-{index}"), metadata: json!({"idx": index}), }, ); envelope.id = format!("msg-{index}"); bus_clone.publish(envelope).expect("publish").sequence })); } let mut published_sequences = Vec::with_capacity(total); for handle in tasks { published_sequences.push(handle.await.expect("join")); } assert_eq!(published_sequences.len(), total); let drained = bus .drain_for_agent("worker", 0) .expect("drain worker inbox should succeed"); assert_eq!(drained.len(), total); for pair in drained.windows(2) { assert!(pair[0].sequence < pair[1].sequence); } let mut seen_tasks = HashSet::new(); for item in drained { if let CoordinationPayload::DelegateTask { task_id, .. } = item.envelope.payload { seen_tasks.insert(task_id); } } assert_eq!(seen_tasks.len(), total); } #[test] fn multi_agent_delegation_flow_updates_context_and_returns_result() { let bus = InMemoryMessageBus::new(); bus.register_agent("lead").expect("register lead"); bus.register_agent("researcher") .expect("register researcher"); let mut request = CoordinationEnvelope::new_direct( "lead", "researcher", "conv-42", "coordination", CoordinationPayload::DelegateTask { task_id: "task-42".to_string(), summary: "Find root cause".to_string(), metadata: json!({"priority": "p1"}), }, ); request.id = "msg-request".to_string(); request.correlation_id = Some("corr-42".to_string()); bus.publish(request.clone()) .expect("request should publish"); let researcher_inbox = bus .drain_for_agent("researcher", 10) .expect("researcher drain"); assert_eq!(researcher_inbox.len(), 1); assert_eq!(researcher_inbox[0].envelope.id, "msg-request"); let mut patch = CoordinationEnvelope::new_broadcast( "researcher", "conv-42", "context", CoordinationPayload::ContextPatch { key: "task-42/findings".to_string(), expected_version: 0, value: json!({"summary": "Root cause isolated"}), }, ); patch.id = "msg-patch".to_string(); patch.correlation_id = Some("corr-42".to_string()); patch.causation_id = Some("msg-request".to_string()); bus.publish(patch).expect("context patch should publish"); let mut result = CoordinationEnvelope::new_direct( "researcher", "lead", "conv-42", "coordination", CoordinationPayload::TaskResult { task_id: "task-42".to_string(), success: true, output: "Investigation complete".to_string(), }, ); result.id = "msg-result".to_string(); result.correlation_id = Some("corr-42".to_string()); result.causation_id = Some("msg-request".to_string()); bus.publish(result).expect("result should publish"); let lead_inbox = bus.drain_for_agent("lead", 10).expect("lead drain"); assert_eq!(lead_inbox.len(), 2); assert_eq!(lead_inbox[0].envelope.id, "msg-patch"); assert_eq!(lead_inbox[1].envelope.id, "msg-result"); let context = bus .context_entry("task-42/findings") .expect("context must exist"); assert_eq!(context.version, 1); assert_eq!(context.updated_by, "researcher"); assert_eq!(context.last_message_id, "msg-patch"); assert_eq!(context.value, json!({"summary": "Root cause isolated"})); } #[test] fn peek_does_not_consume_messages() { let bus = InMemoryMessageBus::new(); bus.register_agent("worker").expect("register worker"); let mut envelope = CoordinationEnvelope::new_direct( "lead", "worker", "conv-peek", "coordination", CoordinationPayload::DelegateTask { task_id: "task-1".to_string(), summary: "peek test".to_string(), metadata: json!({}), }, ); envelope.id = "msg-peek".to_string(); bus.publish(envelope).expect("publish"); let peeked = bus.peek_for_agent("worker", 10).expect("peek"); assert_eq!(peeked.len(), 1); assert_eq!(peeked[0].envelope.id, "msg-peek"); let pending = bus.pending_for_agent("worker").expect("pending"); assert_eq!(pending, 1); } #[test] fn correlation_pending_and_peek_paging_follow_inbox_lifecycle() { let bus = InMemoryMessageBus::new(); bus.register_agent("worker").expect("register worker"); for (message_id, correlation_id) in [ ("msg-corr-0", "corr-a"), ("msg-corr-1", "corr-b"), ("msg-corr-2", "corr-a"), ("msg-corr-3", "corr-a"), ] { let mut envelope = CoordinationEnvelope::new_direct( "lead", "worker", "conv-peek-correlation", "coordination", CoordinationPayload::DelegateTask { task_id: message_id.to_string(), summary: "peek correlation".to_string(), metadata: json!({}), }, ); envelope.id = message_id.to_string(); envelope.correlation_id = Some(correlation_id.to_string()); bus.publish(envelope).expect("publish should succeed"); } assert_eq!( bus.pending_for_agent_correlation("worker", "corr-a") .expect("pending corr-a should succeed"), 3 ); assert_eq!( bus.pending_for_agent_correlation("worker", "corr-b") .expect("pending corr-b should succeed"), 1 ); let page = bus .peek_for_agent_correlation_with_offset("worker", "corr-a", 1, 1) .expect("peek corr-a page should succeed"); assert_eq!(page.len(), 1); assert_eq!(page[0].envelope.id, "msg-corr-2"); let drained_one = bus .drain_for_agent("worker", 1) .expect("drain one should succeed"); assert_eq!(drained_one.len(), 1); assert_eq!(drained_one[0].envelope.id, "msg-corr-0"); assert_eq!( bus.pending_for_agent_correlation("worker", "corr-a") .expect("pending corr-a should succeed after drain"), 2 ); } #[test] fn inbox_correlation_counts_stay_consistent_with_overflow_evictions() { let bus = InMemoryMessageBus::with_limits(InMemoryMessageBusLimits { max_inbox_messages_per_agent: 2, max_dead_letters: 16, max_context_entries: 16, max_seen_message_ids: 32, }); bus.register_agent("worker").expect("register worker"); for (id, corr) in [("m0", "corr-a"), ("m1", "corr-b"), ("m2", "corr-a")] { let mut envelope = CoordinationEnvelope::new_direct( "lead", "worker", "conv-overflow-corr", "coordination", CoordinationPayload::DelegateTask { task_id: id.to_string(), summary: "overflow".to_string(), metadata: json!({}), }, ); envelope.id = id.to_string(); envelope.correlation_id = Some(corr.to_string()); bus.publish(envelope).expect("publish should succeed"); } // m0 (corr-a) should be evicted by inbox overflow. assert_eq!( bus.pending_for_agent_correlation("worker", "corr-a") .expect("corr-a pending should work"), 1 ); assert_eq!( bus.pending_for_agent_correlation("worker", "corr-b") .expect("corr-b pending should work"), 1 ); let corr_a_page = bus .peek_for_agent_correlation_with_offset("worker", "corr-a", 0, 10) .expect("corr-a peek should work"); assert_eq!(corr_a_page.len(), 1); assert_eq!(corr_a_page[0].envelope.id, "m2"); let corr_b_page = bus .peek_for_agent_correlation_with_offset("worker", "corr-b", 0, 10) .expect("corr-b peek should work"); assert_eq!(corr_b_page.len(), 1); assert_eq!(corr_b_page[0].envelope.id, "m1"); } #[test] fn correlation_peek_normalizes_whitespace_in_message_correlation_id() { let bus = InMemoryMessageBus::new(); bus.register_agent("worker").expect("register worker"); let mut envelope = CoordinationEnvelope::new_direct( "lead", "worker", "conv-corr-normalize", "coordination", CoordinationPayload::DelegateTask { task_id: "task-1".to_string(), summary: "normalize".to_string(), metadata: json!({}), }, ); envelope.id = "msg-corr-whitespace".to_string(); envelope.correlation_id = Some(" corr-a ".to_string()); bus.publish(envelope).expect("publish should succeed"); assert_eq!( bus.pending_for_agent_correlation("worker", "corr-a") .expect("pending by normalized correlation should succeed"), 1 ); let page = bus .peek_for_agent_correlation_with_offset("worker", "corr-a", 0, 10) .expect("peek by normalized correlation should succeed"); assert_eq!(page.len(), 1); assert_eq!(page[0].envelope.id, "msg-corr-whitespace"); } #[test] fn registered_agents_and_context_snapshot_are_available() { let bus = InMemoryMessageBus::new(); bus.register_agent("worker-b").expect("register worker-b"); bus.register_agent("worker-a").expect("register worker-a"); let patch = CoordinationEnvelope::new_broadcast( "worker-a", "conv-snapshot", "context", CoordinationPayload::ContextPatch { key: "shared/key".to_string(), expected_version: 0, value: json!({"ok": true}), }, ); bus.publish(patch).expect("publish patch"); let agents = bus.registered_agents(); assert_eq!(agents, vec!["worker-a".to_string(), "worker-b".to_string()]); let snapshot = bus.context_snapshot(); assert_eq!(snapshot.len(), 1); assert_eq!( snapshot .get("shared/key") .expect("shared key should exist") .value, json!({"ok": true}) ); } #[test] fn inbox_limit_drops_oldest_and_records_dead_letter() { let bus = InMemoryMessageBus::with_limits(InMemoryMessageBusLimits { max_inbox_messages_per_agent: 2, max_dead_letters: 8, max_context_entries: 16, max_seen_message_ids: 32, }); bus.register_agent("worker").expect("register worker"); for index in 0..3 { let mut envelope = CoordinationEnvelope::new_direct( "lead", "worker", "conv-limit", "coordination", CoordinationPayload::DelegateTask { task_id: format!("task-{index}"), summary: format!("work-{index}"), metadata: json!({}), }, ); envelope.id = format!("msg-limit-{index}"); envelope.correlation_id = Some("corr-limit".to_string()); bus.publish(envelope).expect("publish should succeed"); } let pending = bus .pending_for_agent("worker") .expect("pending should work"); assert_eq!(pending, 2); assert_eq!( bus.pending_for_agent_correlation("worker", "corr-limit") .expect("pending by correlation should work"), 2 ); let drained = bus.drain_for_agent("worker", 0).expect("drain should work"); assert_eq!(drained.len(), 2); assert_eq!(drained[0].envelope.id, "msg-limit-1"); assert_eq!(drained[1].envelope.id, "msg-limit-2"); assert_eq!( bus.pending_for_agent_correlation("worker", "corr-limit") .expect("pending by correlation after drain should work"), 0 ); let dead_letters = bus.dead_letters(); assert_eq!(dead_letters.len(), 1); assert_eq!(dead_letters[0].envelope.id, "msg-limit-0"); assert!(dead_letters[0].reason.contains("inbox overflow")); let stats = bus.stats(); assert_eq!(stats.publish_attempts_total, 3); assert_eq!(stats.deliveries_total, 3); assert_eq!(stats.inbox_overflow_evictions_total, 1); assert_eq!(stats.dead_letters_total, 1); assert_eq!(stats.dead_letter_evictions_total, 0); assert_eq!(stats.context_evictions_total, 0); assert_eq!(stats.seen_message_id_evictions_total, 0); } #[test] fn dead_letter_limit_is_capped() { let bus = InMemoryMessageBus::with_limits(InMemoryMessageBusLimits { max_inbox_messages_per_agent: 16, max_dead_letters: 2, max_context_entries: 16, max_seen_message_ids: 32, }); bus.register_agent("worker").expect("register worker"); for index in 0..4 { let mut invalid = CoordinationEnvelope::new_direct( "worker", "worker", "conv-dead-letter-limit", "coordination", CoordinationPayload::TaskResult { task_id: format!("task-{index}"), success: false, output: "failed".to_string(), }, ); invalid.id = format!("msg-dead-{index}"); // Missing correlation id -> dead-letter. let _ = bus.publish(invalid); } let dead_letters = bus.dead_letters(); assert_eq!(dead_letters.len(), 2); assert_eq!(dead_letters[0].envelope.id, "msg-dead-2"); assert_eq!(dead_letters[1].envelope.id, "msg-dead-3"); assert_eq!(bus.dead_letter_count(), 2); let stats = bus.stats(); assert_eq!(stats.publish_attempts_total, 0); assert_eq!(stats.deliveries_total, 0); assert_eq!(stats.inbox_overflow_evictions_total, 0); assert_eq!(stats.dead_letters_total, 4); assert_eq!(stats.dead_letter_evictions_total, 2); assert_eq!(stats.context_evictions_total, 0); assert_eq!(stats.seen_message_id_evictions_total, 0); } #[test] fn context_limit_evicts_oldest_entries_and_tracks_stats() { let bus = InMemoryMessageBus::with_limits(InMemoryMessageBusLimits { max_inbox_messages_per_agent: 16, max_dead_letters: 16, max_context_entries: 2, max_seen_message_ids: 32, }); for index in 0..3 { let mut patch = CoordinationEnvelope::new_broadcast( "lead", "conv-context-limit", "delegate.state", CoordinationPayload::ContextPatch { key: format!("delegate/corr-{index}/state"), expected_version: 0, value: json!({"phase":"queued","index":index}), }, ); patch.id = format!("context-msg-{index}"); patch.correlation_id = Some(format!("corr-{index}")); bus.publish(patch).expect("context patch should publish"); } let snapshot = bus.context_snapshot(); assert_eq!(snapshot.len(), 2); assert!(!snapshot.contains_key("delegate/corr-0/state")); assert!(snapshot.contains_key("delegate/corr-1/state")); assert!(snapshot.contains_key("delegate/corr-2/state")); let stats = bus.stats(); assert_eq!(stats.publish_attempts_total, 3); assert_eq!(stats.deliveries_total, 0); assert_eq!(stats.dead_letters_total, 0); assert_eq!(stats.context_evictions_total, 1); assert_eq!(stats.seen_message_id_evictions_total, 0); } #[test] fn context_limit_uses_write_recency_and_preserves_hot_keys() { let bus = InMemoryMessageBus::with_limits(InMemoryMessageBusLimits { max_inbox_messages_per_agent: 16, max_dead_letters: 16, max_context_entries: 2, max_seen_message_ids: 32, }); let mut patch_a = CoordinationEnvelope::new_broadcast( "lead", "conv-context-lru", "delegate.state", CoordinationPayload::ContextPatch { key: "delegate/corr-a/state".to_string(), expected_version: 0, value: json!({"phase":"queued"}), }, ); patch_a.id = "ctx-lru-a0".to_string(); patch_a.correlation_id = Some("corr-a".to_string()); bus.publish(patch_a).expect("first patch should publish"); let mut patch_b = CoordinationEnvelope::new_broadcast( "lead", "conv-context-lru", "delegate.state", CoordinationPayload::ContextPatch { key: "delegate/corr-b/state".to_string(), expected_version: 0, value: json!({"phase":"queued"}), }, ); patch_b.id = "ctx-lru-b0".to_string(); patch_b.correlation_id = Some("corr-b".to_string()); bus.publish(patch_b).expect("second patch should publish"); // Update key A to make it the most recently written key. let mut patch_a_update = CoordinationEnvelope::new_broadcast( "lead", "conv-context-lru", "delegate.state", CoordinationPayload::ContextPatch { key: "delegate/corr-a/state".to_string(), expected_version: 1, value: json!({"phase":"running"}), }, ); patch_a_update.id = "ctx-lru-a1".to_string(); patch_a_update.correlation_id = Some("corr-a".to_string()); bus.publish(patch_a_update) .expect("recency update patch should publish"); let mut patch_c = CoordinationEnvelope::new_broadcast( "lead", "conv-context-lru", "delegate.state", CoordinationPayload::ContextPatch { key: "delegate/corr-c/state".to_string(), expected_version: 0, value: json!({"phase":"queued"}), }, ); patch_c.id = "ctx-lru-c0".to_string(); patch_c.correlation_id = Some("corr-c".to_string()); bus.publish(patch_c) .expect("new key should trigger eviction under limit"); let snapshot = bus.context_snapshot(); assert_eq!(snapshot.len(), 2); assert!(snapshot.contains_key("delegate/corr-a/state")); assert!(snapshot.contains_key("delegate/corr-c/state")); assert!(!snapshot.contains_key("delegate/corr-b/state")); assert_eq!( snapshot .get("delegate/corr-a/state") .expect("A key should remain") .version, 2 ); let stats = bus.stats(); assert_eq!(stats.context_evictions_total, 1); assert_eq!(stats.seen_message_id_evictions_total, 0); } #[test] fn context_entries_recent_with_offset_returns_newest_first_pages() { let bus = InMemoryMessageBus::new(); for key in [ "delegate/corr-a/state", "delegate/corr-b/state", "delegate/corr-c/state", ] { let mut patch = CoordinationEnvelope::new_broadcast( "lead", "conv-context-page", "delegate.state", CoordinationPayload::ContextPatch { key: key.to_string(), expected_version: 0, value: json!({"phase":"queued"}), }, ); patch.id = format!("ctx-page-{key}"); patch.correlation_id = parse_delegate_context_correlation_from_key(key).map(str::to_string); bus.publish(patch).expect("context patch should publish"); } let page = bus.context_entries_recent_with_offset(1, 2); assert_eq!(page.len(), 2); assert_eq!(page[0].0, "delegate/corr-b/state"); assert_eq!(page[1].0, "delegate/corr-a/state"); } #[test] fn dead_letters_recent_returns_newest_first_pages() { let bus = InMemoryMessageBus::new(); bus.register_agent("worker").expect("register worker"); for index in 0..4 { let mut invalid = CoordinationEnvelope::new_direct( "lead", "worker", "conv-dead-letter-page", "delegate.result", CoordinationPayload::TaskResult { task_id: format!("task-{index}"), success: false, output: "failure".to_string(), }, ); invalid.id = format!("dead-page-{index}"); // Missing correlation id causes dead-letter. let _ = bus.publish(invalid); } let page = bus.dead_letters_recent(1, 2); assert_eq!(page.len(), 2); assert_eq!(page[0].envelope.id, "dead-page-2"); assert_eq!(page[1].envelope.id, "dead-page-1"); } #[test] fn context_entries_recent_for_correlation_support_paging_and_count() { let bus = InMemoryMessageBus::new(); let mut patch_a_state = CoordinationEnvelope::new_broadcast( "lead", "conv-correlation-context", "delegate.state", CoordinationPayload::ContextPatch { key: "delegate/corr-a/state".to_string(), expected_version: 0, value: json!({"phase":"queued"}), }, ); patch_a_state.id = "ctx-corr-a-state-0".to_string(); patch_a_state.correlation_id = Some("corr-a".to_string()); bus.publish(patch_a_state) .expect("corr-a state patch should publish"); let mut patch_b_state = CoordinationEnvelope::new_broadcast( "lead", "conv-correlation-context", "delegate.state", CoordinationPayload::ContextPatch { key: "delegate/corr-b/state".to_string(), expected_version: 0, value: json!({"phase":"queued"}), }, ); patch_b_state.id = "ctx-corr-b-state-0".to_string(); patch_b_state.correlation_id = Some("corr-b".to_string()); bus.publish(patch_b_state) .expect("corr-b state patch should publish"); let mut patch_a_state_update = CoordinationEnvelope::new_broadcast( "lead", "conv-correlation-context", "delegate.state", CoordinationPayload::ContextPatch { key: "delegate/corr-a/state".to_string(), expected_version: 1, value: json!({"phase":"running"}), }, ); patch_a_state_update.id = "ctx-corr-a-state-1".to_string(); patch_a_state_update.correlation_id = Some("corr-a".to_string()); bus.publish(patch_a_state_update) .expect("corr-a state update should publish"); let mut patch_a_output = CoordinationEnvelope::new_broadcast( "lead", "conv-correlation-context", "delegate.output", CoordinationPayload::ContextPatch { key: "delegate/corr-a/output".to_string(), expected_version: 0, value: json!({"summary":"done"}), }, ); patch_a_output.id = "ctx-corr-a-output-0".to_string(); patch_a_output.correlation_id = Some("corr-a".to_string()); bus.publish(patch_a_output) .expect("corr-a output patch should publish"); assert_eq!(bus.context_count_for_correlation("corr-a"), 2); assert_eq!(bus.context_count_for_correlation("corr-b"), 1); assert_eq!(bus.context_count_for_correlation("corr-missing"), 0); let page = bus.context_entries_recent_for_correlation_with_offset("corr-a", 0, 2); assert_eq!(page.len(), 2); assert_eq!(page[0].0, "delegate/corr-a/output"); assert_eq!(page[1].0, "delegate/corr-a/state"); let second_page = bus.context_entries_recent_for_correlation_with_offset("corr-a", 1, 1); assert_eq!(second_page.len(), 1); assert_eq!(second_page[0].0, "delegate/corr-a/state"); } #[test] fn delegate_context_indexes_exclude_non_delegate_keys_and_support_paging() { let bus = InMemoryMessageBus::new(); let mut delegate_a_state = CoordinationEnvelope::new_broadcast( "lead", "conv-delegate-context", "delegate.state", CoordinationPayload::ContextPatch { key: "delegate/corr-a/state".to_string(), expected_version: 0, value: json!({"phase":"queued"}), }, ); delegate_a_state.id = "delegate-a-state-0".to_string(); delegate_a_state.correlation_id = Some("corr-a".to_string()); bus.publish(delegate_a_state) .expect("delegate a state patch should publish"); let mut non_delegate = CoordinationEnvelope::new_broadcast( "lead", "conv-delegate-context", "context", CoordinationPayload::ContextPatch { key: "shared/other".to_string(), expected_version: 0, value: json!({"k":"v"}), }, ); non_delegate.id = "shared-other-0".to_string(); non_delegate.correlation_id = Some("corr-a".to_string()); bus.publish(non_delegate) .expect("non-delegate patch should publish"); let mut delegate_a_output = CoordinationEnvelope::new_broadcast( "lead", "conv-delegate-context", "delegate.output", CoordinationPayload::ContextPatch { key: "delegate/corr-a/output".to_string(), expected_version: 0, value: json!({"summary":"done"}), }, ); delegate_a_output.id = "delegate-a-output-0".to_string(); delegate_a_output.correlation_id = Some("corr-a".to_string()); bus.publish(delegate_a_output) .expect("delegate a output patch should publish"); assert_eq!(bus.context_count(), 3); assert_eq!(bus.delegate_context_count(), 2); assert_eq!(bus.delegate_context_count_for_correlation("corr-a"), 2); assert_eq!( bus.delegate_context_count_for_correlation("corr-missing"), 0 ); let all_delegate = bus.delegate_context_entries_recent_with_offset(0, 0); assert_eq!(all_delegate.len(), 2); assert_eq!(all_delegate[0].0, "delegate/corr-a/output"); assert_eq!(all_delegate[1].0, "delegate/corr-a/state"); let delegate_page = bus.delegate_context_entries_recent_for_correlation_with_offset("corr-a", 1, 1); assert_eq!(delegate_page.len(), 1); assert_eq!(delegate_page[0].0, "delegate/corr-a/state"); } #[test] fn dead_letter_correlation_index_tracks_evictions_and_paging() { let bus = InMemoryMessageBus::with_limits(InMemoryMessageBusLimits { max_inbox_messages_per_agent: 16, max_dead_letters: 2, max_context_entries: 16, max_seen_message_ids: 32, }); bus.register_agent("worker").expect("register worker"); let publish_invalid_with_correlation = |message_id: &str, correlation_id: &str| { let mut envelope = CoordinationEnvelope::new_direct( "lead", "missing-worker", "conv-correlation-dead-letters", "delegate.request", CoordinationPayload::DelegateTask { task_id: message_id.to_string(), summary: "should dead-letter".to_string(), metadata: json!({}), }, ); envelope.id = message_id.to_string(); envelope.correlation_id = Some(correlation_id.to_string()); let _ = bus.publish(envelope); }; publish_invalid_with_correlation("dead-corr-a-0", "corr-a"); publish_invalid_with_correlation("dead-corr-b-0", "corr-b"); publish_invalid_with_correlation("dead-corr-a-1", "corr-a"); assert_eq!(bus.dead_letter_count(), 2); assert_eq!(bus.dead_letter_count_for_correlation("corr-a"), 1); assert_eq!(bus.dead_letter_count_for_correlation("corr-b"), 1); assert_eq!(bus.dead_letter_count_for_correlation("corr-missing"), 0); let corr_a_page = bus.dead_letters_recent_for_correlation("corr-a", 0, 2); assert_eq!(corr_a_page.len(), 1); assert_eq!(corr_a_page[0].envelope.id, "dead-corr-a-1"); let corr_a_offset_page = bus.dead_letters_recent_for_correlation("corr-a", 1, 2); assert!(corr_a_offset_page.is_empty()); let corr_b_page = bus.dead_letters_recent_for_correlation("corr-b", 0, 2); assert_eq!(corr_b_page.len(), 1); assert_eq!(corr_b_page[0].envelope.id, "dead-corr-b-0"); } }