#[cfg(feature = "query")]
use crate::agent::clock::{Clock, SystemClock};
use crate::cursor::Cursor;
use crate::error::LaserError;
use crate::laser::Laser;
use crate::provenance::AgentTopic;
use crate::types::{AgentId, MintUlid, PrincipalId};
use laser_wire::agent::{
AgentCard, AgentEnvelope, AgentPresence, ConversationId, Health, OPERATION_CARD,
OPERATION_QUARANTINE, OPERATION_UNQUARANTINE, RecordId,
};
use laser_wire::content::ContentType;
use laser_wire::framing::decode_named;
use std::collections::{HashMap, HashSet, VecDeque};
use std::sync::{Arc, Mutex};
#[cfg(feature = "query")]
const PRESENCE_TTL_MICROS: u64 = 2_000_000;
const APPLIED_FACT_CAPACITY: usize = 16_384;
const MAX_REGISTERED_CARDS: usize = 4_096;
const APPLIED_FACT_TTL_MICROS: u64 = 86_400_000_000;
#[derive(Default)]
pub(crate) struct RegistryCache {
cards: HashMap<AgentId, RegisteredCard>,
quarantined: HashSet<AgentId>,
offsets: Vec<u64>,
presence: HashMap<AgentId, PresenceEntry>,
presence_read_at_micros: Option<u64>,
applied_facts: HashSet<RecordId>,
applied_fact_order: VecDeque<(RecordId, u64)>,
}
#[derive(Clone, Debug)]
pub struct PresenceEntry {
pub presence: AgentPresence,
pub principal: Option<PrincipalId>,
}
#[cfg(feature = "query")]
use laser_wire::clients::{ClientMetadata, ClientMetadataList, ClientMetadataQuery};
#[cfg(feature = "query")]
use laser_wire::codes::{
AGDX_GET_CLIENTS_METADATA_CODE, AGDX_SET_CLIENT_METADATA_CODE, CLIENT_METADATA_OP_VERSION,
};
#[cfg(feature = "query")]
use laser_wire::framing::encode_named;
#[cfg(feature = "query")]
const MAX_PAGE_SIZE: usize = laser_wire::limits::MAX_PAGE_SIZE;
#[derive(Clone, Debug)]
pub struct RegisteredCard {
pub agent: AgentId,
pub card: AgentCard,
pub observed_at_micros: u64,
}
impl RegisteredCard {
pub fn is_fresh(&self, now_micros: u64) -> bool {
match self.card.ttl_micros {
None => true,
Some(ttl) => self.observed_at_micros.saturating_add(ttl) >= now_micros,
}
}
pub fn serves(&self, skill_id: &str) -> bool {
self.card
.capabilities
.iter()
.any(|capability| capability.skill_id == skill_id)
}
pub fn available_for(&self, skill_id: &str) -> bool {
self.card
.capabilities
.iter()
.find(|capability| capability.skill_id == skill_id)
.is_some_and(|capability| capability.health != Some(Health::Unavailable))
}
}
pub struct AgentRegistry<'a> {
#[cfg_attr(not(feature = "query"), allow(dead_code))]
laser: &'a Laser,
cursor: Cursor,
cards: HashMap<AgentId, RegisteredCard>,
presence: HashMap<AgentId, PresenceEntry>,
quarantined: HashSet<AgentId>,
applied_facts: HashSet<RecordId>,
applied_fact_order: VecDeque<(RecordId, u64)>,
cache: Arc<Mutex<RegistryCache>>,
#[cfg_attr(not(feature = "query"), allow(dead_code))]
presence_read_at_micros: Option<u64>,
#[cfg(feature = "sign")]
verifier: Option<Arc<crate::sign::KeyRegistry>>,
}
impl<'a> AgentRegistry<'a> {
pub async fn refresh(&mut self, now_micros: u64) -> Result<usize, LaserError> {
let messages = self.cursor.poll_iggy().await?;
let mut folded = 0;
for (_, message) in &messages {
let Ok(decoded) = crate::agent::consumer::decode_agent_record(
message,
laser_wire::agent::features::NONE,
) else {
continue;
};
let Some(envelope) = decoded.envelope else {
continue;
};
let operation = envelope.operation.as_deref();
let applied = if operation == Some(OPERATION_QUARANTINE)
|| operation == Some(OPERATION_UNQUARANTINE)
{
if envelope
.record
.is_some_and(|id| self.applied_facts.contains(&id))
|| !self.fact_authorized(
&envelope,
decoded.signature_context.as_ref(),
decoded.observed_at_micros,
)
{
false
} else {
if let Some(id) = envelope.record {
self.applied_facts.insert(id);
self.applied_fact_order
.push_back((id, decoded.observed_at_micros));
}
if operation == Some(OPERATION_QUARANTINE) {
apply_quarantine(&mut self.quarantined, &envelope)
} else {
lift_quarantine(&mut self.quarantined, &envelope)
}
}
} else {
apply_card(&mut self.cards, &envelope, decoded.observed_at_micros)
};
if applied {
folded += 1;
}
}
self.cards.retain(|_, card| card.is_fresh(now_micros));
let cutoff = now_micros.saturating_sub(APPLIED_FACT_TTL_MICROS);
while self.applied_fact_order.len() > APPLIED_FACT_CAPACITY
|| self
.applied_fact_order
.front()
.is_some_and(|(_, observed)| *observed < cutoff)
{
if let Some((id, _)) = self.applied_fact_order.pop_front() {
self.applied_facts.remove(&id);
}
}
if let Ok(mut cache) = self.cache.lock() {
cache.cards = self.cards.clone();
cache.quarantined = self.quarantined.clone();
cache.applied_facts = self.applied_facts.clone();
cache.applied_fact_order = self.applied_fact_order.clone();
cache.offsets = self.cursor.offsets().to_vec();
}
Ok(folded)
}
#[cfg(feature = "sign")]
fn fact_authorized(
&self,
envelope: &AgentEnvelope,
context: Option<&laser_wire::agent::SignatureContext>,
observed_at_micros: u64,
) -> bool {
match &self.verifier {
Some(registry) => context.is_some_and(|context| {
registry
.verify_observed_at(envelope, context, observed_at_micros)
.is_ok_and(|verified| verified.kind == crate::sign::KeyKind::Operator)
}),
None => true,
}
}
#[cfg(not(feature = "sign"))]
fn fact_authorized(
&self,
_envelope: &AgentEnvelope,
_context: Option<&laser_wire::agent::SignatureContext>,
_observed_at_micros: u64,
) -> bool {
true
}
pub fn agents(&self) -> impl Iterator<Item = &RegisteredCard> {
self.cards.values()
}
pub fn lookup(&self, agent: &AgentId) -> Option<&RegisteredCard> {
self.cards.get(agent)
}
pub fn resolve(&self, skill_id: &str, now_micros: u64) -> Vec<&RegisteredCard> {
self.cards
.values()
.filter(|card| {
card.available_for(skill_id)
&& card.is_fresh(now_micros)
&& !self.quarantined.contains(&card.agent)
})
.collect()
}
pub fn is_quarantined(&self, agent: &AgentId) -> bool {
self.quarantined.contains(agent)
}
pub fn inbox_for(&self, agent: &AgentId) -> Option<&str> {
self.presence
.get(agent)
.and_then(|entry| entry.presence.inbox.as_deref())
}
pub fn inbox_for_principal(&self, agent: &AgentId, principal: PrincipalId) -> Option<&str> {
self.presence.get(agent).and_then(|entry| {
(entry.principal == Some(principal))
.then_some(entry.presence.inbox.as_deref())
.flatten()
})
}
pub fn principal_for(&self, agent: &AgentId) -> Option<PrincipalId> {
self.presence.get(agent).and_then(|entry| entry.principal)
}
#[cfg(feature = "query")]
pub async fn refresh_presence(&mut self) -> Result<usize, LaserError> {
let now = SystemClock.now_micros();
if let Some(read_at) = self.presence_read_at_micros
&& now.saturating_sub(read_at) < PRESENCE_TTL_MICROS
{
return Ok(self.presence.len());
}
let connections = self
.laser
.client_metadata()
.with_metadata_only(true)
.all()
.await?;
let mut fresh = HashMap::new();
for connection in connections {
if let Some(payload) = connection.metadata.as_deref() {
apply_presence(&mut fresh, payload, connection.user_id);
}
}
let folded = fresh.len();
self.presence = fresh;
self.presence_read_at_micros = Some(now);
if let Ok(mut cache) = self.cache.lock() {
cache.presence = self.presence.clone();
cache.presence_read_at_micros = Some(now);
}
Ok(folded)
}
}
#[cfg(feature = "query")]
fn apply_presence(
presence: &mut HashMap<AgentId, PresenceEntry>,
metadata: &[u8],
user_id: Option<u32>,
) -> bool {
let Ok(body) = decode_named::<AgentPresence>(metadata) else {
return false;
};
if body.validate().is_err() {
return false;
}
let Ok(agent) = AgentId::try_from(body.agent.as_str()) else {
return false;
};
presence.insert(
agent,
PresenceEntry {
presence: body,
principal: user_id.map(PrincipalId::new),
},
);
true
}
impl Laser {
pub fn agent_registry(&self) -> Result<AgentRegistry<'_>, LaserError> {
let cache = self.registry_cache()?;
let (
cards,
quarantined,
applied_facts,
applied_fact_order,
presence,
offsets,
presence_read_at_micros,
) = {
let guard = cache
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
(
guard.cards.clone(),
guard.quarantined.clone(),
guard.applied_facts.clone(),
guard.applied_fact_order.clone(),
guard.presence.clone(),
guard.offsets.clone(),
guard.presence_read_at_micros,
)
};
let cursor = self
.reader(AgentTopic::Registry.topic_string())?
.from_offsets(offsets);
Ok(AgentRegistry {
laser: self,
cursor,
cards,
presence,
quarantined,
applied_facts,
applied_fact_order,
cache,
presence_read_at_micros,
#[cfg(feature = "sign")]
verifier: self.registry_verifier(),
})
}
pub async fn publish_card(&self, source: AgentId, card: &AgentCard) -> Result<(), LaserError> {
let body = laser_wire::framing::encode_named(card)
.map_err(|error| LaserError::Codec(format!("encode card: {error}")))?;
self.agdx(
AgentTopic::Registry,
source.wire_id(),
ConversationId::mint(),
)
.status(OPERATION_CARD)
.body(body)
.content_type(ContentType::Cbor)
.send()
.await?;
Ok(())
}
pub async fn quarantine(&self, operator: AgentId, agent: &AgentId) -> Result<(), LaserError> {
self.publish_registry_fact(OPERATION_QUARANTINE, operator, agent)
.await
}
pub async fn unquarantine(&self, operator: AgentId, agent: &AgentId) -> Result<(), LaserError> {
self.publish_registry_fact(OPERATION_UNQUARANTINE, operator, agent)
.await
}
async fn publish_registry_fact(
&self,
operation: &'static str,
operator: AgentId,
agent: &AgentId,
) -> Result<(), LaserError> {
self.agdx(
AgentTopic::Registry,
operator.wire_id(),
ConversationId::mint(),
)
.status(operation)
.body(agent.as_str().as_bytes().to_vec())
.send()
.await?;
Ok(())
}
#[cfg(feature = "sign")]
pub async fn quarantine_signed(
&self,
operator: AgentId,
agent: &AgentId,
key: &crate::sign::SigningKey,
) -> Result<(), LaserError> {
self.publish_registry_fact_signed(OPERATION_QUARANTINE, operator, agent, key)
.await
}
#[cfg(feature = "sign")]
pub async fn unquarantine_signed(
&self,
operator: AgentId,
agent: &AgentId,
key: &crate::sign::SigningKey,
) -> Result<(), LaserError> {
self.publish_registry_fact_signed(OPERATION_UNQUARANTINE, operator, agent, key)
.await
}
#[cfg(feature = "sign")]
async fn publish_registry_fact_signed(
&self,
operation: &'static str,
operator: AgentId,
agent: &AgentId,
key: &crate::sign::SigningKey,
) -> Result<(), LaserError> {
self.agdx(
AgentTopic::Registry,
operator.wire_id(),
ConversationId::mint(),
)
.status(operation)
.body(agent.as_str().as_bytes().to_vec())
.signed_by(key)
.send()
.await?;
Ok(())
}
#[cfg(feature = "query")]
pub async fn advertise_presence(&self, presence: &AgentPresence) -> Result<(), LaserError> {
let requested = AgentId::try_from(presence.agent.as_str())?;
self.claim_presence(requested)?;
let payload = encode_named(presence)
.map_err(|error| LaserError::Codec(format!("encode presence: {error}")))?;
self.send_raw_with_response(AGDX_SET_CLIENT_METADATA_CODE, payload)
.await?;
Ok(())
}
#[cfg(feature = "query")]
pub async fn clear_presence(&self) -> Result<(), LaserError> {
self.send_raw_with_response(AGDX_SET_CLIENT_METADATA_CODE, Vec::new())
.await?;
self.release_presence();
Ok(())
}
#[cfg(feature = "query")]
pub fn client_metadata(&self) -> ClientMetadataRequest<'_> {
ClientMetadataRequest {
laser: self,
query: ClientMetadataQuery {
v: CLIENT_METADATA_OP_VERSION,
with_metadata_only: false,
user_id: None,
after_client_id: None,
limit: MAX_PAGE_SIZE as u32,
},
}
}
}
#[cfg(feature = "query")]
pub struct ClientMetadataPage {
pub clients: Vec<ClientMetadata>,
pub next_cursor: Option<u32>,
}
#[cfg(feature = "query")]
pub struct ClientMetadataRequest<'a> {
laser: &'a Laser,
query: ClientMetadataQuery,
}
#[cfg(feature = "query")]
impl ClientMetadataRequest<'_> {
pub fn with_metadata_only(mut self, value: bool) -> Self {
self.query.with_metadata_only = value;
self
}
pub fn principal(mut self, principal: PrincipalId) -> Self {
self.query.user_id = Some(principal.get());
self
}
pub fn limit(mut self, limit: u32) -> Self {
self.query.limit = limit;
self
}
pub fn after(mut self, client_id: u32) -> Self {
self.query.after_client_id = Some(client_id);
self
}
pub async fn page(&self) -> Result<ClientMetadataPage, LaserError> {
let payload = encode_named(&self.query)
.map_err(|error| LaserError::Codec(format!("encode request: {error}")))?;
let payload = self
.laser
.send_raw_with_response(AGDX_GET_CLIENTS_METADATA_CODE, payload)
.await?;
let list: ClientMetadataList = decode_named(&payload)?;
Ok(ClientMetadataPage {
clients: list.clients,
next_cursor: list.next_cursor,
})
}
pub async fn all(self) -> Result<Vec<ClientMetadata>, LaserError> {
let mut all = Vec::new();
let mut cursor = self.query.after_client_id;
loop {
let mut request = ClientMetadataRequest {
laser: self.laser,
query: self.query.clone(),
};
request.query.after_client_id = cursor;
let page = request.page().await?;
all.extend(page.clients);
match page.next_cursor {
Some(next) => cursor = Some(next),
None => break,
}
}
Ok(all)
}
}
fn apply_card(
cards: &mut HashMap<AgentId, RegisteredCard>,
envelope: &AgentEnvelope,
observed_at_micros: u64,
) -> bool {
if envelope.operation.as_deref() != Some(OPERATION_CARD) {
return false;
}
let Ok(agent) = AgentId::try_from(envelope.source.as_str()) else {
return false;
};
let Ok(card) = decode_named::<AgentCard>(&envelope.body) else {
return false;
};
if card.validate().is_err() {
return false;
}
if cards.len() >= MAX_REGISTERED_CARDS && !cards.contains_key(&agent) {
let oldest = cards
.iter()
.min_by_key(|(_, registered)| registered.observed_at_micros)
.map(|(id, _)| id.clone());
if let Some(oldest) = oldest {
cards.remove(&oldest);
}
}
cards.insert(
agent.clone(),
RegisteredCard {
agent,
card,
observed_at_micros,
},
);
true
}
fn apply_quarantine(quarantined: &mut HashSet<AgentId>, envelope: &AgentEnvelope) -> bool {
let Ok(body) = std::str::from_utf8(&envelope.body) else {
return false;
};
let Ok(agent) = AgentId::try_from(body) else {
return false;
};
quarantined.insert(agent)
}
fn lift_quarantine(quarantined: &mut HashSet<AgentId>, envelope: &AgentEnvelope) -> bool {
let Ok(body) = std::str::from_utf8(&envelope.body) else {
return false;
};
let Ok(agent) = AgentId::try_from(body) else {
return false;
};
quarantined.remove(&agent)
}
#[cfg(test)]
mod tests {
use super::*;
use laser_wire::agent::{CapabilityDescriptor, ConversationId, Health, RecordId};
use laser_wire::framing::encode_named;
fn descriptor(skill_id: &str) -> CapabilityDescriptor {
CapabilityDescriptor {
skill_id: skill_id.to_owned(),
input: None,
output: None,
cost_class: None,
latency_class: None,
max_concurrency: None,
health: Some(Health::Healthy),
load: None,
}
}
fn card_envelope(source: &str, skills: &[&str], ttl_micros: Option<u64>) -> AgentEnvelope {
let card = AgentCard {
name: None,
version: None,
capabilities: skills.iter().map(|skill| descriptor(skill)).collect(),
ttl_micros,
};
let mut envelope = AgentEnvelope::status(
RecordId::from_u128(1),
ConversationId::from_u128(2),
source.parse().expect("valid agent id"),
OPERATION_CARD,
);
envelope.body = encode_named(&card).expect("card encodes");
envelope
}
#[test]
fn given_cards_when_folded_then_should_keep_the_latest_per_agent_and_resolve_by_capability() {
let mut cards = HashMap::new();
let planner: AgentId = "planner".parse().unwrap();
assert!(apply_card(
&mut cards,
&card_envelope("planner", &["plan"], None),
100
));
assert!(apply_card(
&mut cards,
&card_envelope("worker", &["diagnose"], Some(50)),
100
));
assert_eq!(cards.len(), 2);
assert!(apply_card(
&mut cards,
&card_envelope("planner", &["plan", "summarize"], None),
200
));
assert_eq!(cards.len(), 2);
assert_eq!(cards[&planner].observed_at_micros, 200);
assert!(cards[&planner].serves("summarize"));
let not_a_card = AgentEnvelope::event(
RecordId::from_u128(9),
ConversationId::from_u128(9),
"x".parse().unwrap(),
b"{}".to_vec(),
);
assert!(!apply_card(&mut cards, ¬_a_card, 300));
let worker = cards
.values()
.find(|c| c.agent.as_str() == "worker")
.unwrap();
assert!(worker.is_fresh(150));
assert!(!worker.is_fresh(151));
}
#[test]
fn given_a_quarantine_record_when_folded_then_should_mark_the_named_agent() {
let mut quarantined = HashSet::new();
let worker: AgentId = "worker".parse().unwrap();
let mut envelope = AgentEnvelope::status(
RecordId::from_u128(1),
ConversationId::from_u128(2),
"operator".parse().expect("valid agent id"),
OPERATION_QUARANTINE,
);
envelope.body = b"worker".to_vec();
assert!(apply_quarantine(&mut quarantined, &envelope));
assert!(quarantined.contains(&worker));
assert!(!apply_quarantine(&mut quarantined, &envelope));
let mut bad = envelope.clone();
bad.body = vec![0xff, 0xfe];
assert!(!apply_quarantine(&mut quarantined, &bad));
let mut lift = AgentEnvelope::status(
RecordId::from_u128(3),
ConversationId::from_u128(4),
"operator".parse().expect("valid agent id"),
OPERATION_UNQUARANTINE,
);
lift.body = b"worker".to_vec();
assert!(lift_quarantine(&mut quarantined, &lift));
assert!(!quarantined.contains(&worker));
assert!(!lift_quarantine(&mut quarantined, &lift));
}
#[test]
fn given_advertised_health_when_checking_availability_then_should_skip_only_unavailable() {
let with_health = |health| RegisteredCard {
agent: "a".parse().expect("valid agent id"),
card: AgentCard {
name: None,
version: None,
capabilities: vec![CapabilityDescriptor {
health,
..descriptor("diagnose")
}],
ttl_micros: None,
},
observed_at_micros: 0,
};
assert!(with_health(Some(Health::Healthy)).available_for("diagnose"));
assert!(with_health(Some(Health::Degraded)).available_for("diagnose"));
assert!(with_health(None).available_for("diagnose"));
assert!(
!with_health(Some(Health::Unavailable)).available_for("diagnose"),
"an agent declaring itself unavailable must not be routed to",
);
assert!(!with_health(Some(Health::Healthy)).available_for("other"));
}
#[cfg(feature = "query")]
#[test]
fn given_connection_metadata_when_folded_as_presence_then_should_key_by_agent_and_skip_non_presence()
{
use laser_wire::agent::AgentPresence;
let mut presence = HashMap::new();
let worker: AgentId = "worker".parse().unwrap();
let body =
encode_named(&AgentPresence::new("worker".parse().unwrap()).with_inbox("worker.work"))
.expect("presence encodes");
assert!(apply_presence(&mut presence, &body, Some(7)));
assert_eq!(
presence[&worker].presence.inbox.as_deref(),
Some("worker.work")
);
assert_eq!(presence[&worker].principal, Some(PrincipalId::new(7)));
let moved = encode_named(
&AgentPresence::new("worker".parse().unwrap()).with_inbox("worker.work.v2"),
)
.expect("presence encodes");
assert!(apply_presence(&mut presence, &moved, Some(7)));
assert_eq!(
presence[&worker].presence.inbox.as_deref(),
Some("worker.work.v2")
);
assert!(!apply_presence(&mut presence, b"not a presence body", None));
assert_eq!(presence.len(), 1);
}
}