use crate::agent::budget::BudgetGate;
use crate::agent::clock::{Clock, SystemClock};
use crate::agent::control::ControlBook;
use crate::agent::ctx::AgentCtx;
use crate::agent::pause::{Hold, PauseRuntime, Replay, SourceTopic};
use crate::agent::session::{SessionConfig, Sessions};
use crate::capabilities::HelloOutcome;
use crate::error::LaserError;
use crate::laser::Laser;
use crate::provenance::{AgentTopic, LlmUsage, Provenance};
use crate::stream::{CommitPolicy, Consumer, ConsumerMessage, Headers};
use crate::types::{AgentId, ConsumerGroupName, ConversationId, MessageId};
use async_trait::async_trait;
use futures::FutureExt;
use iggy::prelude::*;
use laser_wire::agent::{
AgentDeadLetter, AgentEnvelope, AgentErrorBody, AgentErrorCode, AgentKind, DeadLetterReason,
LogPosition, OPERATION_TASK, SignatureContext, TaskState, features, validate,
};
use laser_wire::codes::AGENT_OP_VERSION;
use laser_wire::content::ContentType;
use laser_wire::dispatch::{
Dispatch, HandledOperations, addressee_filter, classify, classify_generic,
};
use laser_wire::framing::{decode_named, encode_named};
use laser_wire::headers::{AGENT_VERSION, CONTENT_TYPE, CONVERSATION_ID, FENCE};
use laser_wire::topics::{AGENT_CONTROL, AGENT_SESSIONS};
use std::collections::{BTreeMap, HashSet, VecDeque};
use std::str::FromStr;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::sync::oneshot;
use tokio::time::sleep;
use tracing::{debug, error, warn};
pub(crate) fn backoff_for(attempt: u32) -> Duration {
crate::retry::CONSUMER.delay(attempt)
}
const DEDUP_SCOPE_SEP: char = '\u{1f}';
const FENCE_MAP_SOFT_CAP: usize = 16_384;
#[cfg(feature = "sign")]
const VERIFIED_RECORD_WINDOW: usize = 16_384;
const FENCE_ENTRY_TTL_MICROS: u64 = 600_000_000;
const FENCE_SWEEP_INTERVAL_MICROS: u64 = 1_000_000;
const POLL_INTERVAL: Duration = Duration::from_millis(10);
const DEFAULT_SHUTDOWN_GRACE: Duration = Duration::from_secs(30);
const DEFAULT_MAX_QUEUED_RECORDS: usize = 4_096;
const DEFAULT_MAX_QUEUED_BYTES: usize = 64 * 1024 * 1024;
#[trait_variant::make(AgentHandler: Send)]
pub trait LocalAgentHandler {
async fn handle(&self, message: &AgentMessage, ctx: &AgentCtx<'_>) -> Result<(), LaserError>;
}
#[async_trait]
pub trait AgentMiddleware: Send + Sync {
async fn before_handle(&self, message: &AgentMessage) -> Result<(), LaserError> {
let _ = message;
Ok(())
}
async fn after_handle(
&self,
message: &AgentMessage,
result: &Result<(), LaserError>,
attempt: u32,
) {
let _ = (message, result, attempt);
}
}
#[async_trait]
pub trait DeadLetterSink: Send + Sync {
async fn on_dead_letter(
&self,
message: Option<&AgentMessage>,
capsule: &AgentDeadLetter,
publish_result: &Result<(), LaserError>,
);
}
#[derive(Debug, Clone)]
pub struct AgentMessage {
pub provenance: Provenance,
pub payload: Vec<u8>,
pub id: MessageId,
pub envelope: Option<AgentEnvelope>,
pub content_type: Option<ContentType>,
pub verified_principal: Option<String>,
}
impl AgentMessage {
pub fn body(&self) -> &[u8] {
match &self.envelope {
Some(envelope) => &envelope.body,
None => &self.payload,
}
}
fn from_consumer(
received: &ConsumerMessage,
understood_features: u64,
) -> Result<DecodedAgentMessage, (Box<LaserError>, Vec<u8>)> {
let id = received.position;
let payload = received.payload.to_vec();
if received.headers_malformed {
return Err((
Box::new(LaserError::Invalid(
"the record's header block does not decode".to_owned(),
)),
payload,
));
}
let decoded = match decode_record_parts(
&received.headers,
&received.payload,
received.timestamp_micros,
understood_features,
) {
Ok(decoded) => decoded,
Err(error) => return Err((Box::new(error), payload)),
};
Ok(DecodedAgentMessage {
message: Self {
provenance: decoded.provenance,
payload,
id,
envelope: decoded.envelope,
content_type: decoded.content_type,
verified_principal: None,
},
#[cfg(feature = "sign")]
signature_context: decoded.signature_context,
#[cfg(feature = "sign")]
observed_at_micros: decoded.observed_at_micros,
})
}
}
struct DecodedAgentMessage {
message: AgentMessage,
#[cfg(feature = "sign")]
signature_context: Option<SignatureContext>,
#[cfg(feature = "sign")]
observed_at_micros: u64,
}
pub(crate) struct DecodedAgentRecord {
pub(crate) provenance: Provenance,
pub(crate) envelope: Option<AgentEnvelope>,
pub(crate) content_type: Option<ContentType>,
pub(crate) signature_context: Option<SignatureContext>,
pub(crate) observed_at_micros: u64,
}
pub(crate) fn decode_agent_record(
message: &IggyMessage,
understood_features: u64,
) -> Result<DecodedAgentRecord, LaserError> {
let headers = message.user_headers_map()?.unwrap_or_default();
decode_record_parts(
&headers,
&message.payload,
message.header.timestamp,
understood_features,
)
}
fn decode_record_parts(
headers: &Headers,
payload: &[u8],
observed_at_micros: u64,
understood_features: u64,
) -> Result<DecodedAgentRecord, LaserError> {
let content_type_key = HeaderKey::from_str(CONTENT_TYPE)?;
let content_type_code = headers
.get(&content_type_key)
.map(HeaderValue::as_uint8)
.transpose()?;
let content_type = content_type_code.and_then(ContentType::from_code);
let version_key = HeaderKey::from_str(AGENT_VERSION)?;
let version = headers
.get(&version_key)
.map(HeaderValue::as_uint32)
.transpose()?;
let (provenance, envelope, signature_context) = match version {
Some(version) => {
if version != AGENT_OP_VERSION {
return Err(LaserError::Invalid(format!(
"unsupported agent envelope version {version}"
)));
}
let envelope: AgentEnvelope = decode_named(payload)?;
validate(&envelope)?;
let unmet = envelope.unmet_requirements(understood_features);
if unmet != features::NONE {
return Err(LaserError::Invalid(format!(
"agent envelope requires unsupported features 0x{unmet:016x}"
)));
}
let provenance = provenance_from_envelope(&envelope);
let context = SignatureContext {
content_type: content_type_code,
agent_version: Some(version),
};
(provenance, Some(envelope), Some(context))
}
None => (
crate::provenance::provenance_from_headers(headers)?,
None,
None,
),
};
Ok(DecodedAgentRecord {
provenance,
envelope,
content_type,
signature_context,
observed_at_micros,
})
}
pub(crate) fn provenance_and_envelope(
message: &IggyMessage,
) -> Result<(Provenance, Option<AgentEnvelope>), LaserError> {
let decoded = decode_agent_record(message, features::NONE)?;
Ok((decoded.provenance, decoded.envelope))
}
fn original_conversation(headers: &Headers) -> Option<ConversationId> {
let value = headers.get(&HeaderKey::from_str(CONVERSATION_ID).ok()?)?;
match value.kind() {
HeaderKind::Uint128 => {
Some(laser_wire::agent::ConversationId::from_u128(value.as_uint128().ok()?).into())
}
_ => value.as_str().ok()?.parse().ok(),
}
}
fn provenance_from_envelope(envelope: &AgentEnvelope) -> Provenance {
Provenance::builder()
.conversation_id(envelope.conversation.into())
.maybe_parent_conversation_id(envelope.parent.map(Into::into))
.maybe_root_conversation_id(envelope.root.map(Into::into))
.maybe_agent(AgentId::try_from(envelope.source.as_str()).ok())
.maybe_target_agent_id(
envelope
.target
.as_ref()
.and_then(|target| AgentId::try_from(target.as_str()).ok()),
)
.maybe_idempotency_key(
envelope
.idempotency_key
.as_ref()
.map(|key| key.as_str().to_owned()),
)
.maybe_correlation_id(
envelope
.correlation
.map(|correlation| correlation.to_string()),
)
.maybe_deadline(envelope.deadline_micros.map(IggyTimestamp::from))
.maybe_usage(envelope.usage.map(|usage| LlmUsage {
input_tokens: Some(usage.input_tokens),
output_tokens: Some(usage.output_tokens),
cost_usd: None,
}))
.maybe_fence_token(envelope.metadata.as_ref().and_then(
|metadata| match metadata.get(FENCE) {
Some(laser_wire::query::Value::Uint(fence)) => Some(*fence),
Some(laser_wire::query::Value::Int(fence)) => u64::try_from(*fence).ok(),
_ => None,
},
))
.build()
}
#[derive(Debug, Clone, Copy)]
pub struct RetryPolicy {
pub max_attempts: u32,
pub base_delay: Duration,
}
impl RetryPolicy {
pub fn backoff(max_attempts: u32, base_delay: Duration) -> Self {
Self {
max_attempts,
base_delay,
}
}
fn delay_for(&self, attempt: u32) -> Duration {
self.base_delay
.saturating_mul(2u32.saturating_pow(attempt.min(16)))
}
}
impl Default for RetryPolicy {
fn default() -> Self {
Self {
max_attempts: 5,
base_delay: Duration::from_millis(200),
}
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum ConcurrencyPolicy {
#[default]
Serial,
SerialPerPartition {
max_partitions: usize,
},
}
#[derive(bon::Builder)]
pub struct ReliableConsumer {
pub group: ConsumerGroupName,
pub agent: Option<AgentId>,
#[builder(into)]
pub topic: String,
#[builder(default = 10_000)]
pub dedup_window: usize,
#[builder(default)]
pub retry: RetryPolicy,
#[builder(default = features::NONE)]
pub understood_features: u64,
#[builder(default = POLL_INTERVAL)]
pub poll_interval: Duration,
#[builder(default = DEFAULT_SHUTDOWN_GRACE)]
pub shutdown_grace: Duration,
#[builder(default)]
pub concurrency: ConcurrencyPolicy,
#[builder(default = DEFAULT_MAX_QUEUED_RECORDS)]
pub max_queued_records: usize,
#[builder(default = DEFAULT_MAX_QUEUED_BYTES)]
pub max_queued_bytes: usize,
pub respond_on: Option<AgentTopic<'static>>,
#[builder(default)]
pub inbox_route: crate::agent::router::InboxRoute,
#[builder(default)]
pub ack_on_pickup: bool,
pub deduplicator: Option<Box<dyn Deduplicator>>,
#[builder(default)]
pub warm_dedup: bool,
#[builder(default)]
pub middleware: Vec<std::sync::Arc<dyn AgentMiddleware>>,
pub on_dead_letter: Option<std::sync::Arc<dyn DeadLetterSink>>,
pub operations: Option<Vec<String>>,
pub sessions: Option<SessionConfig>,
#[cfg(feature = "sign")]
pub verifier: Option<std::sync::Arc<crate::sign::KeyRegistry>>,
#[cfg(feature = "sign")]
pub signing_key: Option<std::sync::Arc<crate::sign::SigningKey>>,
}
impl ReliableConsumer {
pub async fn run<H>(
mut self,
laser: &Laser,
handler: H,
ready: oneshot::Sender<()>,
shutdown: oneshot::Receiver<()>,
) -> Result<(), LaserError>
where
H: AgentHandler + Sync + Send + 'static,
{
let stream = laser.stream_required()?.to_owned();
let engine = DeliveryEngine::resolve(laser).await?;
let me = self.agent.as_ref().map(AgentId::wire_id);
if let Some(me) = &me {
engine
.bind_addressee(laser, &stream, &self.topic, self.group.as_str(), me)
.await?;
}
let opener = Opener {
laser: laser.clone(),
stream: stream.clone(),
topic: self.topic.clone(),
group: self.group.as_str().to_owned(),
poll_interval: self.poll_interval,
native: engine.native,
};
let consumer = opener.open().await?;
let deduplicator = self
.deduplicator
.take()
.unwrap_or_else(|| Box::new(SlidingWindow::new(self.dedup_window)));
if self.warm_dedup {
warm_dedup_window(laser, &self, deduplicator.as_ref()).await?;
}
let stream_ident = Identifier::named(&stream)?;
let topic_ident = Identifier::named(&self.topic)?;
let stream_id = laser
.client()
.get_stream(&stream_ident)
.await?
.map(|details| details.id);
let topic_details = laser
.client()
.get_topic(&stream_ident, &topic_ident)
.await?;
let topic_id = topic_details.as_ref().map(|details| details.id);
let topic_generation = topic_details
.as_ref()
.map_or(0, |details| details.created_at.as_micros());
if stream_id.is_none() || topic_id.is_none() {
warn!(
topic = %self.topic,
"could not resolve the numeric stream/topic id, dead-letter capsules \
carry 0 for the unresolved locator half (partition and offset stay correct)"
);
}
let source = match (stream_id, topic_id) {
(Some(stream_id), Some(topic_id)) => Some(SourceTopic {
stream_id,
topic_id,
generation: topic_generation,
}),
_ => None,
};
let (stream_id, topic_id) = (stream_id.unwrap_or_default(), topic_id.unwrap_or_default());
let sessions = laser.sessions_with(self.sessions.clone().unwrap_or_default());
let control = Arc::new(ControlBook::default());
let (reconcile, reconcile_requests) = tokio::sync::mpsc::unbounded_channel();
let mut pause = None;
let mut late_control = false;
let mut control_follower = if self.topic == AGENT_CONTROL {
None
} else {
match crate::agent::pause::load_control(laser, me.as_ref(), &control).await? {
Some(feed) => {
match (&me, source) {
(Some(me), Some(source)) => {
pause = Some(Arc::new(PauseRuntime::new(
sessions.clone(),
me.clone(),
Arc::clone(&control),
&feed,
source,
reconcile.clone(),
)));
}
(Some(_), None) => {
warn!(topic = %self.topic, "the source topic ids did not resolve, sessions of this agent cannot be paused");
}
(None, _) => {}
}
control.set_live(true);
let (stop, stopped) = oneshot::channel();
let task = tokio::spawn(crate::agent::pause::follow_control(
laser.clone(),
feed,
me.clone(),
Arc::clone(&control),
pause.is_some().then(|| reconcile.clone()),
self.poll_interval,
stopped,
));
Some((stop, task))
}
None => {
late_control = true;
None
}
}
};
drop(reconcile);
let probe_interval = sessions.config().heartbeat_value();
let reliable = Arc::new(ReliableWorker {
handler,
laser: laser.clone(),
retry: self.retry,
understood_features: self.understood_features,
dedup: deduplicator,
agent: self.agent,
respond_on: self.respond_on,
inbox_route: self.inbox_route,
ack_on_pickup: self.ack_on_pickup,
stream_id,
topic_id,
middleware: self.middleware,
on_dead_letter: self.on_dead_letter,
topic: self.topic.clone(),
operations: self.operations,
sessions,
control,
pause: Arc::new(
pause
.clone()
.map_or_else(std::sync::OnceLock::new, std::sync::OnceLock::from),
),
budget: BudgetGate::default(),
high_water_fence: dashmap::DashMap::new(),
fence_last_sweep: std::sync::atomic::AtomicU64::new(0),
#[cfg(feature = "sign")]
verifier: self.verifier,
#[cfg(feature = "sign")]
signing_key: self.signing_key,
#[cfg(feature = "sign")]
verified_records: Mutex::new(DedupWindow::new(VERIFIED_RECORD_WINDOW)),
});
let pause_driver = match &pause {
Some(runtime) => {
if let Err(error) = runtime.recover().await {
stop_follower(control_follower).await;
return Err(error);
}
let (stop, stopped) = oneshot::channel();
let worker: Arc<dyn Replay> = reliable.clone();
let task = tokio::spawn(crate::agent::pause::drive(
Arc::clone(runtime),
worker,
reconcile_requests,
stopped,
));
Some((stop, task))
}
None => None,
};
let late_driver: LateDriver = Arc::default();
if late_control {
let start_pause = match (&me, source) {
(Some(me), Some(source)) => {
let slot = Arc::clone(&reliable.pause);
let sessions = reliable.sessions.clone();
let me = me.clone();
let book = Arc::clone(&reliable.control);
let worker: Arc<dyn Replay> = reliable.clone();
let driver = Arc::clone(&late_driver);
let start: crate::agent::pause::StartPause = Box::new(move |feed| {
let (reconcile, requests) = tokio::sync::mpsc::unbounded_channel();
let runtime = Arc::new(PauseRuntime::new(
sessions,
me,
book,
feed,
source,
reconcile.clone(),
));
Box::pin(async move {
start_late_pause(runtime, worker, requests, reconcile, slot, driver)
.await
})
});
Some(start)
}
_ => None,
};
let (stop, stopped) = oneshot::channel();
let task = tokio::spawn(crate::agent::pause::discover_then_follow_control(
laser.clone(),
me.clone(),
Arc::clone(&reliable.control),
start_pause,
self.poll_interval,
stopped,
));
control_follower = Some((stop, task));
}
let _late = AbortLateOnDrop(Arc::clone(&late_driver));
let _background = AbortOnDrop(
control_follower
.iter()
.map(|(_, task)| task.abort_handle())
.chain(pause_driver.iter().map(|(_, task)| task.abort_handle()))
.collect(),
);
let _ = ready.send(());
let stop = async {
if shutdown.await.is_err() {
std::future::pending::<()>().await;
}
}
.fuse();
let result = match self.concurrency {
ConcurrencyPolicy::Serial => {
run_serial(&opener, consumer, &reliable, stop, self.shutdown_grace).await
}
ConcurrencyPolicy::SerialPerPartition { max_partitions } => {
let limits = LaneLimits {
max_partitions: max_partitions.max(1),
max_queued_records: self.max_queued_records.clamp(1, u32::MAX as usize),
max_queued_bytes: self.max_queued_bytes.clamp(1, u32::MAX as usize),
probe_interval,
};
run_per_partition(
&opener,
consumer,
Arc::clone(&reliable),
limits,
stop,
self.shutdown_grace,
)
.await
}
};
if let Some((stop, mut task)) = pause_driver {
let _ = stop.send(());
if tokio::time::timeout(self.shutdown_grace, &mut task)
.await
.is_err()
{
warn!("the pause driver did not finish within the shutdown grace, aborting it");
task.abort();
}
}
let late = late_driver
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.take();
if let Some((stop, mut task)) = late {
let _ = stop.send(());
if tokio::time::timeout(self.shutdown_grace, &mut task)
.await
.is_err()
{
task.abort();
}
}
stop_follower(control_follower).await;
result
}
}
type LateDriver = Arc<Mutex<Option<(oneshot::Sender<()>, tokio::task::JoinHandle<()>)>>>;
struct AbortLateOnDrop(LateDriver);
impl Drop for AbortLateOnDrop {
fn drop(&mut self) {
if let Some((_, task)) = self
.0
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.take()
{
task.abort();
}
}
}
async fn start_late_pause(
runtime: Arc<PauseRuntime>,
worker: Arc<dyn Replay>,
requests: tokio::sync::mpsc::UnboundedReceiver<ConversationId>,
reconcile: tokio::sync::mpsc::UnboundedSender<ConversationId>,
slot: Arc<std::sync::OnceLock<Arc<PauseRuntime>>>,
driver: LateDriver,
) -> Option<tokio::sync::mpsc::UnboundedSender<ConversationId>> {
let mut attempt = 0;
loop {
match runtime.recover().await {
Ok(()) => break,
Err(error) => {
attempt += 1;
if attempt >= crate::retry::CONSUMER_FAILURE_BUDGET {
warn!(%error, "the pause runtime for the late control topic did not recover, sessions of this agent cannot be paused");
return None;
}
tokio::time::sleep(crate::retry::CONSUMER.delay(attempt)).await;
}
}
}
if slot.set(Arc::clone(&runtime)).is_err() {
return None;
}
let (stop, stopped) = oneshot::channel();
let task = tokio::spawn(crate::agent::pause::drive(
runtime, worker, requests, stopped,
));
*driver
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = Some((stop, task));
Some(reconcile)
}
struct AbortOnDrop(Vec<tokio::task::AbortHandle>);
impl Drop for AbortOnDrop {
fn drop(&mut self) {
for task in &self.0 {
task.abort();
}
}
}
async fn stop_follower(follower: Option<(oneshot::Sender<()>, tokio::task::JoinHandle<()>)>) {
if let Some((stop, task)) = follower {
let _ = stop.send(());
let _ = task.await;
}
}
#[derive(Clone, Copy)]
struct DeliveryEngine {
native: bool,
filters: bool,
}
impl DeliveryEngine {
async fn resolve(laser: &Laser) -> Result<Self, LaserError> {
let mut capabilities = laser.capabilities().await;
if capabilities.hello == HelloOutcome::Unknown {
capabilities = laser.refresh_capabilities().await;
}
if !crate::stream::transport::policy_aware(&capabilities)? {
return Ok(Self {
native: true,
filters: false,
});
}
if laser.connection_string().is_none() {
if capabilities.managed {
tracing::warn!(
"the agent reads natively because its client was not built from a connection string, so no group filter is bound"
);
}
return Ok(Self {
native: true,
filters: false,
});
}
Ok(Self {
native: false,
filters: capabilities.filters.native && capabilities.filters.catalog,
})
}
async fn bind_addressee(
&self,
laser: &Laser,
stream: &str,
topic: &str,
group: &str,
me: &laser_wire::agent::AgentId,
) -> Result<(), LaserError> {
if !self.filters || (topic != AGENT_SESSIONS && topic != AGENT_CONTROL) {
return Ok(());
}
laser
.stream(stream)
.topic(topic)
.consumer_group(group)
.create()
.filter(addressee_filter(me))
.build()
.await
.map(|_| ())
}
}
#[derive(Clone)]
struct Opener {
laser: Laser,
stream: String,
topic: String,
group: String,
poll_interval: Duration,
native: bool,
}
impl Opener {
async fn open(&self) -> Result<Consumer, LaserError> {
let builder = self
.laser
.stream(&self.stream)
.topic(&self.topic)
.consumer_group(&self.group)
.consumer()
.commit_policy(CommitPolicy::Disabled)
.poll_interval(self.poll_interval)
.create_group(true);
let builder = if self.native {
builder.native()
} else {
builder
};
builder.build().await
}
async fn reopen(&self, consumer: &mut Consumer, failures: &mut u32) -> Result<(), LaserError> {
if let Err(error) = consumer.shutdown().await {
debug!(%error, topic = %self.topic, "the failed consumer did not leave its group cleanly");
}
loop {
sleep(backoff_for(*failures)).await;
match self.open().await {
Ok(reopened) => {
*consumer = reopened;
return Ok(());
}
Err(error) if recoverable(&error, *failures + 1) => {
*failures += 1;
warn!(%error, attempt = *failures, topic = %self.topic, "reopening the agent consumer failed, retrying");
}
Err(error) => return Err(error),
}
}
}
}
pub(crate) const MAX_CONSECUTIVE_POLL_ERRORS: u32 = crate::retry::CONSUMER_FAILURE_BUDGET;
fn recoverable(error: &LaserError, failures: u32) -> bool {
let rejoin = matches!(
error,
LaserError::Iggy(IggyError::ConsumerGroupMemberNotFound(..))
);
(rejoin || error.is_retryable()) && failures < MAX_CONSECUTIVE_POLL_ERRORS
}
async fn run_serial<H, S>(
opener: &Opener,
mut consumer: Consumer,
worker: &ReliableWorker<H>,
stop: S,
shutdown_grace: Duration,
) -> Result<(), LaserError>
where
H: AgentHandler + Sync + Send + 'static,
S: std::future::Future<Output = ()> + futures::future::FusedFuture,
{
tokio::pin!(stop);
let mut failures = 0u32;
let mut stopping = false;
let result = loop {
if stopping {
break Ok(());
}
let next = tokio::select! {
() = &mut stop => break Ok(()),
next = consumer.next() => next,
};
let message = match next {
None => break Ok(()),
Some(Ok(message)) => message,
Some(Err(error)) => {
failures += 1;
if !recoverable(&error, failures) {
break Err(error);
}
warn!(%error, attempt = failures, topic = %opener.topic, "agent consumer failed, reopening it");
if let Err(error) = opener.reopen(&mut consumer, &mut failures).await {
break Err(error);
}
worker.invalidate_pause();
continue;
}
};
failures = 0;
worker.observe_assignment(&consumer);
let handled = worker.consume(&message);
tokio::pin!(handled);
let outcome = tokio::select! {
outcome = &mut handled => outcome,
() = &mut stop => {
stopping = true;
match tokio::time::timeout(shutdown_grace, &mut handled).await {
Ok(outcome) => outcome,
Err(_) => return Err(LaserError::Timeout("agent shutdown drain")),
}
}
};
if let Err(error) = outcome {
break Err(error);
}
if let Err(error) = consumer.commit(&message).await {
failures += 1;
if !recoverable(&error, failures) {
break Err(error);
}
warn!(%error, source = %message.position, "committing a handled record failed, reopening the consumer");
if let Err(error) = opener.reopen(&mut consumer, &mut failures).await {
break Err(error);
}
worker.invalidate_pause();
}
};
if let Err(error) = consumer.shutdown().await {
warn!(%error, "failed to leave the consumer group on shutdown");
}
result
}
struct LaneLimits {
max_partitions: usize,
max_queued_records: usize,
max_queued_bytes: usize,
probe_interval: Duration,
}
async fn run_per_partition<H, S>(
opener: &Opener,
mut consumer: Consumer,
worker: Arc<ReliableWorker<H>>,
limits: LaneLimits,
stop: S,
shutdown_grace: Duration,
) -> Result<(), LaserError>
where
H: AgentHandler + Sync + Send + 'static,
S: std::future::Future<Output = ()> + futures::future::FusedFuture,
{
tokio::pin!(stop);
let mut failures = 0u32;
loop {
let (outcome, drain) =
drive_lanes(opener, &mut consumer, &worker, &limits, &mut stop).await;
if drain.delivered {
failures = 0;
}
let tolerate_commits = matches!(outcome, LaneOutcome::Failed(_));
let drained =
tokio::time::timeout(shutdown_grace, drain.finish(&consumer, tolerate_commits)).await;
match outcome {
LaneOutcome::Stopped => {
return match drained {
Ok(Ok(())) => {
if let Err(error) = consumer.shutdown().await {
warn!(%error, "failed to leave the consumer group on shutdown");
}
Ok(())
}
Ok(Err(error)) => Err(error),
Err(_) => Err(LaserError::Timeout("agent shutdown drain")),
};
}
LaneOutcome::Fatal(error) => {
if matches!(drained, Ok(Ok(())))
&& let Err(error) = consumer.shutdown().await
{
warn!(%error, "failed to leave the consumer group on shutdown");
}
return Err(error);
}
LaneOutcome::Failed(error) => {
if let Ok(Err(drain_error)) | Err(drain_error) =
drained.map_err(|_| LaserError::Timeout("agent lane drain"))
{
return Err(drain_error);
}
failures += 1;
if !recoverable(&error, failures) {
return Err(error);
}
warn!(%error, attempt = failures, topic = %opener.topic, "agent consumer failed, reopening it");
opener.reopen(&mut consumer, &mut failures).await?;
worker.invalidate_pause();
}
}
}
}
enum LaneOutcome {
Stopped,
Failed(LaserError),
Fatal(LaserError),
}
struct Lane {
sender: tokio::sync::mpsc::Sender<QueuedMessage>,
task: tokio::task::JoinHandle<()>,
}
struct LaneDrain {
lanes: Vec<tokio::task::JoinHandle<()>>,
completions: tokio::sync::mpsc::Receiver<Result<ConsumerMessage, LaserError>>,
delivered: bool,
}
impl LaneDrain {
async fn finish(
mut self,
consumer: &Consumer,
tolerate_commits: bool,
) -> Result<(), LaserError> {
for lane in self.lanes {
let _ = lane.await;
}
while let Some(completion) = self.completions.recv().await {
let message = completion?;
match consumer.commit(&message).await {
Ok(()) => {}
Err(error) if tolerate_commits => {
debug!(%error, source = %message.position, "a handled record could not be committed, it is read again");
}
Err(error) => return Err(error),
}
}
Ok(())
}
}
async fn drive_lanes<H, S>(
opener: &Opener,
consumer: &mut Consumer,
worker: &Arc<ReliableWorker<H>>,
limits: &LaneLimits,
stop: &mut std::pin::Pin<&mut S>,
) -> (LaneOutcome, LaneDrain)
where
H: AgentHandler + Sync + Send + 'static,
S: std::future::Future<Output = ()> + futures::future::FusedFuture,
{
let mut lanes: std::collections::HashMap<u32, Lane> = std::collections::HashMap::new();
let mut retiring: std::collections::HashMap<u32, tokio::task::JoinHandle<()>> =
std::collections::HashMap::new();
let (commit_tx, mut commit_rx) = tokio::sync::mpsc::channel::<
Result<ConsumerMessage, LaserError>,
>(limits.max_partitions * 2);
let notify = Arc::new(tokio::sync::Notify::new());
let record_capacity = Arc::new(tokio::sync::Semaphore::new(limits.max_queued_records));
let byte_capacity = Arc::new(tokio::sync::Semaphore::new(limits.max_queued_bytes));
let mut probe = tokio::time::interval(limits.probe_interval);
probe.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
probe.reset();
let mut delivered = false;
let outcome = 'polling: loop {
while let Ok(completion) = commit_rx.try_recv() {
match completion {
Ok(message) => {
if let Err(error) = consumer.commit_handled(&message).await {
break 'polling LaneOutcome::Failed(error);
}
}
Err(error) => break 'polling LaneOutcome::Fatal(error),
}
}
tokio::select! {
() = &mut *stop => break LaneOutcome::Stopped,
() = notify.notified() => {}
_ = probe.tick() => {
drop_revoked(opener, consumer, &mut lanes, &mut retiring).await;
}
next = consumer.next() => match next {
Some(Ok(message)) => {
delivered = true;
worker.observe_assignment(consumer);
let partition = message.partition_id;
if !lanes.contains_key(&partition) && lanes.len() >= limits.max_partitions {
if let Err(error) = worker.consume(&message).await {
break 'polling LaneOutcome::Fatal(error);
}
if let Err(error) = consumer.commit_handled(&message).await {
break 'polling LaneOutcome::Failed(error);
}
continue;
}
let lane = lanes.entry(partition).or_insert_with(|| {
spawn_lane(
worker,
retiring.remove(&partition),
limits.max_queued_records,
commit_tx.clone(),
Arc::clone(¬ify),
)
});
let buffered_bytes = message.payload.len()
+ message.user_headers.as_ref().map_or(0, bytes::Bytes::len);
let Ok(record_permit) = Arc::clone(&record_capacity).acquire_owned().await
else {
break 'polling LaneOutcome::Fatal(LaserError::HandlerConfig(
"record queue closed".to_owned(),
));
};
let byte_permits = buffered_bytes.clamp(1, limits.max_queued_bytes) as u32;
let Ok(byte_permit) = Arc::clone(&byte_capacity)
.acquire_many_owned(byte_permits)
.await
else {
break 'polling LaneOutcome::Fatal(LaserError::HandlerConfig(
"byte queue closed".to_owned(),
));
};
let queued = QueuedMessage {
message,
record_permit,
byte_permit,
};
if lane.sender.send(queued).await.is_err() {
continue;
}
}
Some(Err(error)) => {
if matches!(error, LaserError::Iggy(IggyError::ConsumerGroupMemberNotFound(..))) {
drop_revoked(opener, consumer, &mut lanes, &mut retiring).await;
}
break LaneOutcome::Failed(error);
}
None => break LaneOutcome::Stopped,
},
}
};
let mut tasks: Vec<_> = lanes.into_values().map(|lane| lane.task).collect();
tasks.extend(retiring.into_values());
drop(commit_tx);
(
outcome,
LaneDrain {
lanes: tasks,
completions: commit_rx,
delivered,
},
)
}
fn spawn_lane<H>(
worker: &Arc<ReliableWorker<H>>,
predecessor: Option<tokio::task::JoinHandle<()>>,
capacity: usize,
commits: tokio::sync::mpsc::Sender<Result<ConsumerMessage, LaserError>>,
notify: Arc<tokio::sync::Notify>,
) -> Lane
where
H: AgentHandler + Sync + Send + 'static,
{
let (sender, mut queue) = tokio::sync::mpsc::channel::<QueuedMessage>(capacity);
let worker = Arc::clone(worker);
let task = tokio::spawn(async move {
if let Some(predecessor) = predecessor {
let _ = predecessor.await;
}
while let Some(queued) = queue.recv().await {
let result = worker.consume(&queued.message).await;
drop(queued.record_permit);
drop(queued.byte_permit);
let failed = result.is_err();
let _ = commits.send(result.map(|()| queued.message)).await;
notify.notify_one();
if failed {
break;
}
}
});
Lane { sender, task }
}
async fn drop_revoked(
opener: &Opener,
consumer: &Consumer,
lanes: &mut std::collections::HashMap<u32, Lane>,
retiring: &mut std::collections::HashMap<u32, tokio::task::JoinHandle<()>>,
) {
let assigned = match consumer.assigned_partitions() {
Some(assigned) => Some(assigned),
None => opener.assigned_partitions().await,
};
let Some(assigned) = assigned else {
return;
};
let revoked: Vec<u32> = lanes
.keys()
.filter(|partition| !assigned.contains(partition))
.copied()
.collect();
for partition in revoked {
if let Some(lane) = lanes.remove(&partition) {
debug!(
partition,
"dropping the lane of a partition no longer assigned"
);
drop(lane.sender);
retiring.insert(partition, lane.task);
}
}
retiring.retain(|_, task| !task.is_finished());
}
impl Opener {
async fn assigned_partitions(&self) -> Option<std::collections::BTreeSet<u32>> {
let client = self.laser.client();
let me = match client.get_me().await {
Ok(me) => me.client_id,
Err(error) => {
debug!(%error, "the assignment probe could not identify this client");
return None;
}
};
let group = client
.get_consumer_group(
&Identifier::named(&self.stream).ok()?,
&Identifier::named(&self.topic).ok()?,
&Identifier::named(&self.group).ok()?,
)
.await;
match group {
Ok(Some(details)) => details
.members
.iter()
.find(|member| member.id == me)
.map(|member| member.partitions.iter().copied().collect()),
Ok(None) => None,
Err(error) => {
debug!(%error, "the assignment probe could not read the group");
None
}
}
}
}
struct QueuedMessage {
message: ConsumerMessage,
record_permit: tokio::sync::OwnedSemaphorePermit,
byte_permit: tokio::sync::OwnedSemaphorePermit,
}
fn dedup_key(provenance: &Provenance) -> Option<String> {
let key = provenance.idempotency_key.as_ref()?;
Some(match &provenance.agent {
Some(agent) => format!("{}{DEDUP_SCOPE_SEP}{key}", agent.as_str()),
None => key.clone(),
})
}
#[derive(Clone, Copy)]
struct FenceEntry {
fence: u64,
touched_micros: u64,
}
fn accept_fence(
high_water: &dashmap::DashMap<ConversationId, FenceEntry>,
last_sweep_micros: &std::sync::atomic::AtomicU64,
task: ConversationId,
fence: u64,
now_micros: u64,
) -> bool {
if high_water.len() > FENCE_MAP_SOFT_CAP {
use std::sync::atomic::Ordering;
let previous = last_sweep_micros.load(Ordering::Relaxed);
if now_micros.saturating_sub(previous) > FENCE_SWEEP_INTERVAL_MICROS
&& last_sweep_micros
.compare_exchange(previous, now_micros, Ordering::Relaxed, Ordering::Relaxed)
.is_ok()
{
high_water.retain(|_, entry| {
now_micros.saturating_sub(entry.touched_micros) < FENCE_ENTRY_TTL_MICROS
});
}
}
let mut entry = high_water.entry(task).or_insert(FenceEntry {
fence: 0,
touched_micros: now_micros,
});
if fence < entry.fence {
return false;
}
entry.fence = fence;
entry.touched_micros = now_micros;
true
}
const ANONYMOUS_AGENT: &str = "anonymous";
struct ReliableWorker<H> {
handler: H,
laser: Laser,
retry: RetryPolicy,
understood_features: u64,
dedup: Box<dyn Deduplicator>,
agent: Option<AgentId>,
respond_on: Option<AgentTopic<'static>>,
inbox_route: crate::agent::router::InboxRoute,
ack_on_pickup: bool,
stream_id: u32,
topic_id: u32,
middleware: Vec<std::sync::Arc<dyn AgentMiddleware>>,
on_dead_letter: Option<std::sync::Arc<dyn DeadLetterSink>>,
topic: String,
operations: Option<Vec<String>>,
sessions: Sessions,
control: Arc<ControlBook>,
pause: Arc<std::sync::OnceLock<Arc<PauseRuntime>>>,
budget: BudgetGate,
#[cfg(feature = "sign")]
verifier: Option<std::sync::Arc<crate::sign::KeyRegistry>>,
#[cfg(feature = "sign")]
signing_key: Option<std::sync::Arc<crate::sign::SigningKey>>,
high_water_fence: dashmap::DashMap<ConversationId, FenceEntry>,
fence_last_sweep: std::sync::atomic::AtomicU64,
#[cfg(feature = "sign")]
verified_records: Mutex<DedupWindow>,
}
impl<H> ReliableWorker<H> {
fn observe_assignment(&self, consumer: &Consumer) {
if let Some(pause) = self.pause.get() {
pause.observe_assignment(consumer.assigned_partitions());
}
}
fn invalidate_pause(&self) {
if let Some(pause) = self.pause.get() {
pause.invalidate();
}
}
fn dispatch(&self, message: &AgentMessage) -> Dispatch {
let me = match (&self.agent, &message.provenance.target_agent_id) {
(Some(agent), _) | (None, Some(agent)) => agent.wire_id(),
(None, None) => ANONYMOUS_AGENT.parse().expect("a valid agent id"),
};
match &message.envelope {
Some(envelope) => {
let operations: Vec<&str> = self
.operations
.iter()
.flatten()
.map(String::as_str)
.collect();
let handled = match self.operations {
Some(_) => HandledOperations::Only(&operations),
None => HandledOperations::Any,
};
classify(envelope, &self.topic, &me, handled)
}
None => classify_generic(
message
.provenance
.target_agent_id
.as_ref()
.map(AgentId::as_str),
message.provenance.causal_parent.is_some(),
message.provenance.correlation_id.is_some(),
&self.topic,
&me,
),
}
}
async fn pick_up_submitted(
&self,
message: &AgentMessage,
) -> Option<crate::agent::SessionLease> {
let envelope = message.envelope.as_ref()?;
let submitted = envelope
.metadata
.as_ref()
.and_then(|metadata| metadata.get(laser_wire::agent::METADATA_SUBMITTED))
.is_some_and(|value| matches!(value, laser_wire::query::Value::Bool(true)));
let agent = self.agent.as_ref()?;
if !submitted || envelope.kind != AgentKind::Command {
return None;
}
let session = self
.sessions
.open(message.provenance.conversation_id)
.as_agent(agent.wire_id());
match session.pick_up().await {
Ok(lease) => Some(lease),
Err(error) => {
warn!(%error, source = %message.id, "failed to mark the submitted session working");
None
}
}
}
fn log_position(&self, id: MessageId) -> LogPosition {
LogPosition::new(self.stream_id, self.topic_id, id.partition_id, id.offset)
}
async fn dead_letter(
&self,
message: &AgentMessage,
reason: DeadLetterReason,
attempts: u32,
detail: &str,
) -> Result<(), LaserError> {
let capsule = AgentDeadLetter {
source: self.log_position(message.id),
reason,
attempts,
detail: Some(detail.to_owned()),
payload: message.payload.clone(),
};
let mut provenance = message.provenance.clone();
provenance.causal_parent = Some(message.id);
provenance.deadline = None;
let session = Some(message.provenance.conversation_id);
self.publish_dead_letter(provenance, message.id, capsule, Some(message), session)
.await
}
async fn dead_letter_undecodable(
&self,
source: MessageId,
conversation: Option<ConversationId>,
payload: Vec<u8>,
) -> Result<(), LaserError> {
let capsule = AgentDeadLetter {
source: self.log_position(source),
reason: DeadLetterReason::DecodeFailed,
attempts: 0,
detail: None,
payload,
};
let position = capsule.source;
let provenance = Provenance::builder()
.conversation_id(conversation.unwrap_or_else(|| {
ConversationId::derive(&format!(
"dead-letter\u{1f}{}\u{1f}{}\u{1f}{}\u{1f}{}",
position.stream_id, position.topic_id, position.partition_id, position.offset
))
}))
.causal_parent(source)
.build();
self.publish_dead_letter(provenance, source, capsule, None, conversation)
.await
}
async fn publish_dead_letter(
&self,
provenance: Provenance,
source: MessageId,
capsule: AgentDeadLetter,
message: Option<&AgentMessage>,
session: Option<ConversationId>,
) -> Result<(), LaserError> {
let reason = capsule.reason;
let result = self.send_dead_letter(&provenance, &capsule).await;
if let Err(error) = &result {
error!(%error, source = %source, ?reason, "failed to publish the dead-letter capsule, leaving the source offset uncommitted");
}
if let Some(sink) = &self.on_dead_letter {
sink.on_dead_letter(message, &capsule, &result).await;
}
if result.is_ok()
&& let Some(session) = session
{
self.fail_session_on_dead_letter(session, &capsule).await;
}
result
}
async fn fail_session_on_dead_letter(
&self,
session: ConversationId,
capsule: &AgentDeadLetter,
) {
if !self.sessions.config().fails_on_dead_letter() {
return;
}
let Some(agent) = &self.agent else {
warn!(%session, "a consumer without an agent id cannot fail the session of a dead letter");
return;
};
let source = capsule.source;
let body = AgentErrorBody {
code: AgentErrorCode::Internal,
message: Some(match capsule.detail.as_deref() {
Some(detail) => format!(
"a record of this session was dead-lettered ({:?}): {detail}",
capsule.reason
),
None => format!(
"a record of this session was dead-lettered ({:?})",
capsule.reason
),
}),
retryable: false,
detail: Some(BTreeMap::from([
(
"dead_letter_reason".to_owned(),
laser_wire::query::Value::Uint(u64::from(capsule.reason.code())),
),
(
"source".to_owned(),
laser_wire::query::Value::Str(format!(
"{}/{}/{}/{}",
source.stream_id, source.topic_id, source.partition_id, source.offset
)),
),
])),
};
if let Err(error) = self
.sessions
.open(session)
.as_agent(agent.wire_id())
.fail(body)
.await
{
warn!(%error, %session, "failed to fail the session of a dead-lettered record");
}
}
async fn send_dead_letter(
&self,
provenance: &Provenance,
capsule: &AgentDeadLetter,
) -> Result<(), LaserError> {
let payload = encode_named(capsule)
.map_err(|error| LaserError::Codec(format!("dead-letter capsule: {error}")))?;
let mut headers = BTreeMap::<HeaderKey, HeaderValue>::try_from(provenance)
.map_err(|error| LaserError::Codec(format!("dead-letter headers: {error}")))?;
let key = HeaderKey::from_str(CONTENT_TYPE)?;
headers.insert(key, HeaderValue::from(ContentType::Cbor.code()));
let topic = AgentTopic::Dlq.topic_string();
let partition_key = provenance.partition_key();
self.laser
.send_with_headers(&topic, payload, headers, Some(&partition_key))
.await
.map(|_| ())
}
}
impl<H> ReliableWorker<H>
where
H: AgentHandler + Sync,
{
#[tracing::instrument(target = "laser", level = "debug", skip_all, fields(conversation = tracing::field::Empty, operation = "handle"))]
async fn consume(&self, received: &ConsumerMessage) -> Result<(), LaserError> {
self.deliver(received, true).await
}
async fn deliver(&self, received: &ConsumerMessage, live: bool) -> Result<(), LaserError> {
let source = received.position;
let decoded = match AgentMessage::from_consumer(received, self.understood_features) {
Ok(decoded) => decoded,
Err((error, payload)) => {
warn!(%error, source = %source, "undecodable provenance, dead-lettering raw payload");
let conversation = (!received.headers_malformed)
.then(|| original_conversation(&received.headers))
.flatten();
self.dead_letter_undecodable(source, conversation, payload)
.await?;
return Ok(());
}
};
#[cfg(feature = "sign")]
let DecodedAgentMessage {
mut message,
signature_context,
observed_at_micros,
} = decoded;
#[cfg(not(feature = "sign"))]
let DecodedAgentMessage { message, .. } = decoded;
tracing::Span::current().record(
"conversation",
tracing::field::display(&message.provenance.conversation_id),
);
if let (Some(target), Some(agent)) = (&message.provenance.target_agent_id, &self.agent)
&& target != agent
{
debug!(target = %target, agent = %agent, source = %message.id, "skipping message targeted at another agent");
return Ok(());
}
let dispatch = self.dispatch(&message);
if dispatch != Dispatch::Work {
debug!(source = %message.id, dispatch = dispatch.as_str(), "skipping a record that is not work for this agent");
return Ok(());
}
#[cfg(feature = "sign")]
if let Some(registry) = &self.verifier {
let verified = message.envelope.as_ref().and_then(|envelope| {
signature_context.as_ref().and_then(|context| {
registry
.verify_observed_at(envelope, context, observed_at_micros)
.ok()
})
});
let Some(verified) = verified else {
warn!(source = %message.id, "unsigned or unverified message on a verified topic, dead-lettering");
self.dead_letter(
&message,
DeadLetterReason::Rejected,
0,
"signature verification failed",
)
.await?;
return Ok(());
};
if let Some(record) = message
.envelope
.as_ref()
.and_then(|envelope| envelope.record)
{
let first_delivery = self
.verified_records
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.insert(&record.to_string());
if !first_delivery {
warn!(source = %message.id, %record, "replayed signed record, dropping");
return Ok(());
}
}
message.verified_principal = Some(verified.principal);
}
let _gate = match (self.pause.get(), live) {
(Some(pause), true) => {
let gate = pause.gate(message.provenance.conversation_id).await;
if pause.hold(self, &message).await? == Hold::Commit {
return Ok(());
}
Some(gate)
}
_ => None,
};
if let Some(agent) = &self.agent {
let session = message.provenance.conversation_id;
let admitted = self
.budget
.admit(
&self.sessions,
session,
received.partition_id,
received.current_offset,
|| {
let mut lens = self.sessions.open(session).as_agent(agent.wire_id());
lens.control = Some(Arc::clone(&self.control));
lens
},
)
.await;
if !admitted {
debug!(source = %message.id, "skipping work for a session over its budget");
return Ok(());
}
}
if let Some(fence) = message.provenance.fence_token
&& !accept_fence(
&self.high_water_fence,
&self.fence_last_sweep,
message.provenance.conversation_id,
fence,
SystemClock.now_micros(),
)
{
debug!(source = %message.id, fence, "stale-holder fence, dropping replay");
return Ok(());
}
if let Some(key) = dedup_key(&message.provenance) {
if !self.dedup.observe(&key).await {
debug!(dedup_key = %key, source = %message.id, "skipping duplicate message");
return Ok(());
}
}
if let Some(deadline) = message.provenance.deadline
&& IggyTimestamp::now().as_micros() > deadline.as_micros()
{
warn!(source = %message.id, "message past its deadline, dead-lettering");
self.dead_letter(
&message,
DeadLetterReason::DeadlineExceeded,
0,
"message past its deadline",
)
.await?;
return Ok(());
}
if self.ack_on_pickup
&& let (Some(agent), Some(respond_on), Some(envelope)) =
(&self.agent, &self.respond_on, &message.envelope)
&& envelope.kind == AgentKind::Command
&& let Some(correlation) = envelope.correlation
{
let producer =
self.laser
.agdx(respond_on.clone(), agent.wire_id(), envelope.conversation);
let ack = producer
.status(OPERATION_TASK)
.with_correlation(correlation)
.with_task_state(TaskState::Working);
#[cfg(feature = "sign")]
let ack = match &self.signing_key {
Some(signing_key) => ack.signed_by(signing_key),
None => ack,
};
if let Err(error) = ack.send().await {
warn!(source = %message.id, %error, "failed to emit ack-on-pickup status");
}
}
let _pickup = self.pick_up_submitted(&message).await;
let ctx = AgentCtx::new(
&self.laser,
&message,
self.agent.clone(),
self.respond_on.clone(),
self.inbox_route.clone(),
#[cfg(feature = "sign")]
self.signing_key.clone(),
)
.with_request_at(self.log_position(message.id))
.with_sessions(self.sessions.clone(), Arc::clone(&self.control));
for middleware in &self.middleware {
if let Err(error) = middleware.before_handle(&message).await {
warn!(%error, source = %message.id, "middleware rejected message before handling, dead-lettering");
self.dead_letter(&message, DeadLetterReason::Rejected, 0, &error.to_string())
.await?;
return Ok(());
}
}
let mut attempt = 0;
loop {
let result = std::panic::AssertUnwindSafe(self.handler.handle(&message, &ctx))
.catch_unwind()
.await
.unwrap_or_else(|panic| {
let reason = panic
.downcast_ref::<&str>()
.map(|message| (*message).to_owned())
.or_else(|| panic.downcast_ref::<String>().cloned())
.unwrap_or_else(|| "unknown panic".to_owned());
Err(LaserError::HandlerConfig(format!(
"handler panicked: {reason}"
)))
});
for middleware in &self.middleware {
middleware
.after_handle(&message, &result, attempt + 1)
.await;
}
match result {
Ok(()) => {
debug!(source = %message.id, "message handled");
return Ok(());
}
Err(error) => {
if !error.is_retryable() {
warn!(%error, source = %message.id, "handler rejected message, dead-lettering without retry");
self.dead_letter(
&message,
DeadLetterReason::Rejected,
attempt + 1,
&error.to_string(),
)
.await?;
return Ok(());
}
if attempt + 1 >= self.retry.max_attempts {
error!(%error, source = %message.id, attempts = attempt + 1, "handler exhausted retries, dead-lettering");
self.dead_letter(
&message,
DeadLetterReason::RetryExhausted,
attempt + 1,
&error.to_string(),
)
.await?;
return Ok(());
}
warn!(%error, source = %message.id, attempt = attempt + 1, "handler failed, retrying");
sleep(self.retry.delay_for(attempt)).await;
attempt += 1;
}
}
}
}
}
impl<H> Replay for ReliableWorker<H>
where
H: AgentHandler + Sync,
{
fn replay<'a>(
&'a self,
received: &'a ConsumerMessage,
) -> futures::future::BoxFuture<'a, Result<(), LaserError>> {
Box::pin(self.deliver(received, false))
}
}
async fn warm_dedup_window(
laser: &Laser,
settings: &ReliableConsumer,
dedup: &dyn Deduplicator,
) -> Result<(), LaserError> {
let stream = Identifier::named(laser.stream_required()?)?;
let topic_id = Identifier::named(&settings.topic)?;
let Some(details) = laser.client().get_topic(&stream, &topic_id).await? else {
return Ok(());
};
let group_consumer =
iggy::prelude::Consumer::group(Identifier::named(settings.group.as_str())?);
let reader = iggy::prelude::Consumer::new(Identifier::named("laser-dedup-warmer")?);
let depth = u64::try_from(settings.dedup_window).unwrap_or(u64::MAX);
for partition in 0..crate::poll::bounded_partitions(details.partitions_count) {
let Some(offset) = laser
.client()
.get_consumer_offset(&group_consumer, &stream, &topic_id, Some(partition))
.await?
else {
continue;
};
let stored = offset.stored_offset;
let start = stored.saturating_sub(depth.saturating_sub(1));
let count =
u32::try_from(stored.saturating_sub(start).saturating_add(1)).unwrap_or(u32::MAX);
let polled = laser
.client()
.poll_messages(
&stream,
&topic_id,
Some(partition),
&reader,
&PollingStrategy::offset(start),
count,
false,
)
.await?;
for message in polled.messages {
if message.header.offset > stored {
continue;
}
if let Some(key) = warm_dedup_key(
&message,
settings.agent.as_ref(),
settings.understood_features,
#[cfg(feature = "sign")]
settings.verifier.as_deref(),
) {
dedup.observe(&key).await;
}
}
}
Ok(())
}
fn warm_dedup_key(
message: &IggyMessage,
agent: Option<&AgentId>,
understood_features: u64,
#[cfg(feature = "sign")] verifier: Option<&crate::sign::KeyRegistry>,
) -> Option<String> {
let decoded = decode_agent_record(message, understood_features).ok()?;
if let (Some(target), Some(agent)) = (&decoded.provenance.target_agent_id, agent)
&& target != agent
{
return None;
}
#[cfg(feature = "sign")]
if let Some(registry) = verifier {
registry
.verify_observed_at(
decoded.envelope.as_ref()?,
decoded.signature_context.as_ref()?,
decoded.observed_at_micros,
)
.ok()?;
}
dedup_key(&decoded.provenance)
}
#[async_trait]
pub trait Deduplicator: Send + Sync {
async fn observe(&self, key: &str) -> bool;
}
pub struct SlidingWindow {
inner: Mutex<DedupWindow>,
}
impl SlidingWindow {
pub fn new(capacity: usize) -> Self {
Self {
inner: Mutex::new(DedupWindow::new(capacity)),
}
}
}
#[async_trait]
impl Deduplicator for SlidingWindow {
async fn observe(&self, key: &str) -> bool {
self.inner
.lock()
.expect("the dedup mutex is not poisoned")
.insert(key)
}
}
struct DedupWindow {
capacity: usize,
seen: HashSet<String>,
order: VecDeque<String>,
}
impl DedupWindow {
fn new(capacity: usize) -> Self {
Self {
capacity: capacity.max(1),
seen: HashSet::new(),
order: VecDeque::new(),
}
}
fn insert(&mut self, key: &str) -> bool {
if self.seen.contains(key) {
return false;
}
if self.order.len() >= self.capacity
&& let Some(evicted) = self.order.pop_front()
{
self.seen.remove(&evicted);
}
self.seen.insert(key.to_owned());
self.order.push_back(key.to_owned());
true
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn given_consumer_failures_when_backing_off_then_should_jitter_within_the_upper_half() {
for attempt in [1, 3, 10] {
let ceiling = crate::retry::CONSUMER.ceiling(attempt);
for _ in 0..64 {
let delay = backoff_for(attempt);
assert!(
delay <= ceiling && delay >= ceiling / 2,
"attempt {attempt}: {delay:?}"
);
}
}
assert_eq!(crate::retry::CONSUMER.ceiling(1), Duration::from_millis(50));
assert_eq!(crate::retry::CONSUMER.ceiling(40), Duration::from_secs(1));
}
#[test]
fn given_an_agdx_command_when_warming_dedup_then_should_restore_the_scoped_key() {
let envelope = AgentEnvelope::command(
laser_wire::agent::RecordId::from_u128(3),
laser_wire::agent::ConversationId::from_u128(1),
"author".parse().expect("source parses"),
laser_wire::agent::CorrelationId::from_u128(2),
b"work".to_vec(),
)
.with_operation(laser_wire::agent::OPERATION_CHAT)
.with_target("worker".parse().expect("target parses"))
.with_idempotency_key("attempt-1".parse().expect("key parses"));
let mut headers = BTreeMap::new();
headers.insert(
HeaderKey::from_str(AGENT_VERSION).expect("version key"),
HeaderValue::from(AGENT_OP_VERSION),
);
let message = IggyMessage::builder()
.payload(bytes::Bytes::from(
encode_named(&envelope).expect("envelope encodes"),
))
.user_headers(headers)
.build()
.expect("message builds");
let worker: AgentId = "worker".parse().expect("worker parses");
let other: AgentId = "other".parse().expect("other parses");
assert_eq!(
warm_dedup_key(
&message,
Some(&worker),
features::NONE,
#[cfg(feature = "sign")]
None,
),
Some(format!("author{DEDUP_SCOPE_SEP}attempt-1"))
);
assert!(
warm_dedup_key(
&message,
Some(&other),
features::NONE,
#[cfg(feature = "sign")]
None,
)
.is_none()
);
}
#[test]
fn given_a_small_positive_fence_when_decoded_then_should_preserve_the_token() {
let envelope = AgentEnvelope::command(
laser_wire::agent::RecordId::from_u128(3),
laser_wire::agent::ConversationId::from_u128(1),
"orchestrator".parse().expect("agent id parses"),
laser_wire::agent::CorrelationId::from_u128(2),
Vec::new(),
)
.with_metadata(FENCE, laser_wire::query::Value::Int(7));
assert_eq!(provenance_from_envelope(&envelope).fence_token, Some(7));
}
#[test]
fn given_an_envelope_with_usage_when_decoded_then_should_keep_token_counts() {
let envelope = AgentEnvelope::command(
laser_wire::agent::RecordId::from_u128(3),
laser_wire::agent::ConversationId::from_u128(1),
"author".parse().expect("agent id parses"),
laser_wire::agent::CorrelationId::from_u128(2),
b"work".to_vec(),
)
.with_operation(laser_wire::agent::OPERATION_CHAT)
.with_usage(laser_wire::agent::TokenUsage {
input_tokens: 13,
output_tokens: 21,
..Default::default()
});
let usage = provenance_from_envelope(&envelope)
.usage
.expect("usage should be preserved");
assert_eq!(usage.input_tokens, Some(13));
assert_eq!(usage.output_tokens, Some(21));
}
#[test]
fn given_a_seen_key_when_inserting_again_then_should_report_a_duplicate() {
let mut window = DedupWindow::new(8);
assert!(window.insert("a"));
assert!(!window.insert("a"));
assert!(window.insert("b"));
}
#[test]
fn given_a_full_window_when_inserting_then_should_evict_the_oldest_key() {
let mut window = DedupWindow::new(2);
assert!(window.insert("a"));
assert!(window.insert("b"));
assert!(window.insert("c"));
assert!(window.insert("a"));
}
#[test]
fn given_increasing_attempts_when_computing_backoff_then_should_grow_and_stay_bounded() {
let policy = RetryPolicy::backoff(5, Duration::from_millis(100));
assert_eq!(policy.delay_for(0), Duration::from_millis(100));
assert_eq!(policy.delay_for(1), Duration::from_millis(200));
assert_eq!(policy.delay_for(2), Duration::from_millis(400));
assert!(policy.delay_for(60) >= policy.delay_for(2));
}
#[test]
fn given_two_agents_with_the_same_idempotency_key_when_scoped_then_should_differ() {
let conversation = ConversationId::new();
let with_agent = |agent: &str| {
Provenance::builder()
.conversation_id(conversation)
.agent(agent.parse().expect("valid agent id"))
.idempotency_key("attempt-1".to_owned())
.build()
};
assert_ne!(
dedup_key(&with_agent("planner")),
dedup_key(&with_agent("worker"))
);
let anon = Provenance::builder()
.conversation_id(conversation)
.idempotency_key("attempt-1".to_owned())
.build();
assert_eq!(dedup_key(&anon).as_deref(), Some("attempt-1"));
}
#[test]
fn given_a_fence_gate_when_a_stale_token_arrives_then_should_drop_it_and_keep_per_task_scope() {
let high_water = dashmap::DashMap::new();
let sweep = std::sync::atomic::AtomicU64::new(0);
let task_a = ConversationId::new();
let task_b = ConversationId::new();
assert!(accept_fence(&high_water, &sweep, task_a, 1, 100));
assert!(accept_fence(&high_water, &sweep, task_a, 2, 101));
assert!(!accept_fence(&high_water, &sweep, task_a, 1, 102));
assert!(accept_fence(&high_water, &sweep, task_a, 2, 103));
assert!(accept_fence(&high_water, &sweep, task_b, 1, 104));
}
}