use crate::error::LaserError;
use crate::govern::{ActionCounters, ActionKind, GovernedAction};
use crate::laser::Laser;
use crate::provenance::AgentTopic;
use crate::types::MintUlid;
use iggy::prelude::{HeaderKey, HeaderValue};
use laser_wire::agent::{
AgentEnvelope, AgentErrorBody, AgentId, AgentKind, ChannelId, ConversationId, CorrelationId,
IdempotencyKey, LogPosition, RecordId, TaskState, TokenUsage, validate,
};
use laser_wire::codes::AGENT_OP_VERSION;
use laser_wire::content::ContentType;
use laser_wire::framing::{decode_named, encode_named};
use laser_wire::headers::{AGENT_VERSION, CONTENT_TYPE, CONVERSATION_ID, TARGET_AGENT_ID};
use laser_wire::query::Value;
use std::collections::BTreeMap;
use std::str::FromStr;
use std::time::Duration;
pub const DEFAULT_CHUNK_FLUSH_BYTES: usize = 512;
pub const DEFAULT_CHUNK_LINGER_MS: u64 = 20;
pub const MAX_CHUNK_BODY_BYTES: usize = 64 * 1024;
#[derive(Clone)]
pub struct Agdx {
laser: Laser,
topic: String,
source: AgentId,
conversation: ConversationId,
}
impl Laser {
pub fn agdx(
&self,
topic: AgentTopic<'_>,
source: AgentId,
conversation: ConversationId,
) -> Agdx {
Agdx {
laser: self.clone(),
topic: topic.topic_string(),
source,
conversation,
}
}
}
impl Agdx {
pub fn command(&self, correlation: CorrelationId, body: Vec<u8>) -> AgdxSend<'_> {
self.send_of(AgentEnvelope::command(
RecordId::mint(),
self.conversation,
self.source.clone(),
correlation,
body,
))
}
pub fn respond(&self, correlation: CorrelationId, body: Vec<u8>) -> AgdxSend<'_> {
self.send_of(AgentEnvelope::response(
RecordId::mint(),
self.conversation,
self.source.clone(),
correlation,
body,
))
}
pub fn emit(&self, body: Vec<u8>) -> AgdxSend<'_> {
self.send_of(AgentEnvelope::event(
RecordId::mint(),
self.conversation,
self.source.clone(),
body,
))
}
pub fn status(&self, operation: impl Into<String>) -> AgdxSend<'_> {
self.send_of(AgentEnvelope::status(
RecordId::mint(),
self.conversation,
self.source.clone(),
operation,
))
}
pub fn fail(
&self,
correlation: CorrelationId,
error: &AgentErrorBody,
) -> Result<AgdxSend<'_>, LaserError> {
let body = encode_named(error)?;
Ok(self
.send_of(AgentEnvelope::error(
RecordId::mint(),
self.conversation,
self.source.clone(),
correlation,
body,
))
.content_type(ContentType::Cbor))
}
pub fn stream(&self, correlation: CorrelationId, purpose: impl Into<String>) -> AgdxStream {
AgdxStream {
agdx: self.clone(),
correlation,
channel: ChannelId::mint(),
purpose: purpose.into(),
sequence: 0,
deadline_micros: None,
target: None,
content_type: ContentType::Raw,
buffer: None,
}
}
pub async fn request_input(
&self,
reply_topic: AgentTopic<'_>,
prompt: impl Into<Vec<u8>>,
timeout: Duration,
) -> Result<Vec<u8>, LaserError> {
let interrupt = CorrelationId::mint();
let mut reader = self.laser.agdx_reply_reader(reply_topic).await?;
self.command(interrupt, prompt.into()).send().await?;
let reply = self
.laser
.await_agdx_reply(&mut reader, interrupt, timeout)
.await?;
if reply.kind == AgentKind::Error {
let message = decode_named::<AgentErrorBody>(&reply.body)
.ok()
.and_then(|body| body.message)
.unwrap_or_else(|| "the input request was rejected".to_owned());
return Err(LaserError::Rejected(message));
}
Ok(reply.body)
}
fn send_of(&self, envelope: AgentEnvelope) -> AgdxSend<'_> {
AgdxSend {
agdx: self,
envelope,
content_type: ContentType::Raw,
#[cfg(feature = "sign")]
sign_key: None,
claim_check: None,
}
}
#[tracing::instrument(target = "laser", level = "debug", skip_all, fields(conversation = %self.conversation, correlation = envelope.correlation.as_ref().map(tracing::field::display), topic = %self.topic, operation = "agdx"))]
async fn publish(
&self,
envelope: AgentEnvelope,
content_type: ContentType,
) -> Result<Option<RecordId>, LaserError> {
let record = envelope.record;
let (message, partition_key) = self.assemble(envelope, content_type)?;
self.laser
.send_batch(&self.topic, vec![message], Some(&partition_key))
.await?;
Ok(record)
}
fn assemble(
&self,
envelope: AgentEnvelope,
content_type: ContentType,
) -> Result<(iggy::prelude::IggyMessage, String), LaserError> {
validate(&envelope)?;
let payload = encode_named(&envelope)?;
let headers = agdx_headers(&envelope, content_type)?;
let partition_key = envelope.conversation.to_string();
let message = iggy::prelude::IggyMessage::builder()
.payload(bytes::Bytes::from(payload))
.user_headers(headers)
.build()?;
Ok((message, partition_key))
}
}
#[must_use = "an unsent envelope does nothing until you call .send().await"]
pub struct AgdxSend<'a> {
agdx: &'a Agdx,
envelope: AgentEnvelope,
content_type: ContentType,
#[cfg(feature = "sign")]
sign_key: Option<&'a crate::sign::SigningKey>,
claim_check: Option<(&'a dyn crate::blob::BlobStore, usize)>,
}
impl<'a> AgdxSend<'a> {
pub fn with_target(mut self, target: AgentId) -> Self {
self.envelope = self.envelope.with_target(target);
self
}
pub fn with_cause(mut self, cause: RecordId, cause_at: Option<LogPosition>) -> Self {
self.envelope = self.envelope.with_cause(cause, cause_at);
self
}
pub fn with_correlation(mut self, correlation: CorrelationId) -> Self {
self.envelope = self.envelope.with_correlation(correlation);
self
}
pub fn with_idempotency_key(mut self, key: IdempotencyKey) -> Self {
self.envelope = self.envelope.with_idempotency_key(key);
self
}
pub fn with_deadline_micros(mut self, deadline_micros: u64) -> Self {
self.envelope = self.envelope.with_deadline_micros(deadline_micros);
self
}
pub fn with_task_state(mut self, state: TaskState) -> Self {
self.envelope = self.envelope.with_task_state(state);
self
}
pub fn with_operation(mut self, operation: impl Into<String>) -> Self {
self.envelope = self.envelope.with_operation(operation);
self
}
pub fn with_tool(mut self, tool: impl Into<String>) -> Self {
self.envelope = self.envelope.with_tool(tool);
self
}
pub fn with_usage(mut self, usage: TokenUsage) -> Self {
self.envelope = self.envelope.with_usage(usage);
self
}
pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<Value>) -> Self {
self.envelope = self.envelope.with_metadata(key, value);
self
}
pub fn last(mut self) -> Self {
self.envelope.last = true;
self
}
pub fn content_type(mut self, content_type: ContentType) -> Self {
self.content_type = content_type;
self
}
pub fn body(mut self, body: impl Into<Vec<u8>>) -> Self {
self.envelope.body = body.into();
self
}
#[cfg(feature = "sign")]
pub fn signed_by(mut self, key: &'a crate::sign::SigningKey) -> Self {
self.sign_key = Some(key);
self
}
pub fn claim_check(
mut self,
store: &'a dyn crate::blob::BlobStore,
threshold_bytes: usize,
) -> Self {
self.claim_check = Some((store, threshold_bytes));
self
}
pub async fn send(self) -> Result<Option<RecordId>, LaserError> {
let mut envelope = self.envelope;
let mut content_type = self.content_type;
#[cfg(feature = "sign")]
let will_sign = self.sign_key.is_some();
#[cfg(not(feature = "sign"))]
let will_sign = false;
let correlation = envelope.correlation.map(|value| value.to_string());
let action = GovernedAction {
kind: action_kind(envelope.kind),
stream: self.agdx.laser.stream_required()?,
topic: &self.agdx.topic,
source: Some(envelope.source.as_str()),
target: envelope.target.as_ref().map(|target| target.as_str()),
conversation: Some(envelope.conversation.into()),
correlation: correlation.as_deref(),
operation: envelope.operation.as_deref(),
tool: envelope.tool.as_deref(),
on_behalf_of: metadata_str(&envelope, laser_wire::agent::METADATA_DELEGATED_BY),
purpose: metadata_str(&envelope, laser_wire::agent::METADATA_PURPOSE),
data_classification: metadata_str(
&envelope,
laser_wire::agent::METADATA_DATA_CLASSIFICATION,
),
payload: &envelope.body,
signed: will_sign,
counters: ActionCounters::default(),
};
if let Some(modified) = self.agdx.laser.govern(action).await? {
envelope.body = modified;
}
if let Some((store, threshold)) = self.claim_check {
let body = std::mem::take(&mut envelope.body);
let (body, replaced) = crate::blob::check_in(store, threshold, body).await?;
envelope.body = body;
if let Some(replaced) = replaced {
content_type = replaced;
}
}
#[cfg(feature = "sign")]
let envelope = if let Some(key) = self.sign_key {
let context = laser_wire::agent::SignatureContext {
content_type: Some(content_type.code()),
agent_version: Some(AGENT_OP_VERSION),
};
let signature = key.sign_with_context(&envelope, context)?;
envelope.with_signature(signature)
} else {
envelope
};
self.agdx.publish(envelope, content_type).await
}
}
pub struct AgdxStream {
agdx: Agdx,
correlation: CorrelationId,
channel: ChannelId,
purpose: String,
sequence: u64,
deadline_micros: Option<u64>,
target: Option<AgentId>,
content_type: ContentType,
buffer: Option<ChunkBuffer>,
}
struct ChunkBuffer {
max_chunks: usize,
linger: std::time::Duration,
first_at: Option<std::time::Instant>,
messages: Vec<iggy::prelude::IggyMessage>,
}
impl AgdxStream {
pub fn channel(&self) -> ChannelId {
self.channel
}
pub fn with_deadline_micros(mut self, deadline_micros: u64) -> Self {
self.deadline_micros = deadline_micros.into();
self
}
pub fn with_target(mut self, target: AgentId) -> Self {
self.target = Some(target);
self
}
pub fn content_type(mut self, content_type: ContentType) -> Self {
self.content_type = content_type;
self
}
#[must_use]
pub fn buffered(mut self, max_chunks: usize, linger: std::time::Duration) -> Self {
self.buffer = Some(ChunkBuffer {
max_chunks: max_chunks.max(1),
linger,
first_at: None,
messages: Vec::new(),
});
self
}
pub async fn write(&mut self, body: Vec<u8>) -> Result<(), LaserError> {
ensure_chunk_body_within_cap(&body)?;
let envelope = self.chunk(body, false, None, None);
self.sequence += 1;
match &mut self.buffer {
None => {
self.agdx.publish(envelope, self.content_type).await?;
}
Some(buffer) => {
let (message, _) = self.agdx.assemble(envelope, self.content_type)?;
buffer.first_at.get_or_insert_with(std::time::Instant::now);
buffer.messages.push(message);
let linger_tripped = buffer
.first_at
.is_some_and(|first| first.elapsed() >= buffer.linger);
if buffer.messages.len() >= buffer.max_chunks || linger_tripped {
self.flush().await?;
}
}
}
Ok(())
}
pub async fn flush(&mut self) -> Result<(), LaserError> {
let Some(buffer) = &mut self.buffer else {
return Ok(());
};
if buffer.messages.is_empty() {
return Ok(());
}
let batch = std::mem::take(&mut buffer.messages);
buffer.first_at = None;
let partition_key = self.agdx.conversation.to_string();
self.agdx
.laser
.send_batch(&self.agdx.topic, batch, Some(&partition_key))
.await
.map(|_| ())
}
pub async fn finish(
mut self,
finish_reason: impl Into<String>,
usage: Option<TokenUsage>,
) -> Result<(), LaserError> {
let envelope = self.chunk(Vec::new(), true, Some(finish_reason.into()), usage);
self.send_terminal(envelope, self.content_type).await
}
pub async fn fail(mut self, error: &AgentErrorBody) -> Result<(), LaserError> {
let body = encode_named(error)?;
let mut envelope = AgentEnvelope::error(
RecordId::mint(),
self.agdx.conversation,
self.agdx.source.clone(),
self.correlation,
body,
);
envelope.channel = Some(self.channel);
envelope.sequence = Some(self.sequence);
if let Some(target) = self.target.clone() {
envelope = envelope.with_target(target);
}
self.send_terminal(envelope, ContentType::Cbor).await
}
async fn send_terminal(
&mut self,
envelope: AgentEnvelope,
content_type: ContentType,
) -> Result<(), LaserError> {
match &mut self.buffer {
None => {
self.agdx.publish(envelope, content_type).await?;
}
Some(buffer) => {
let (message, partition_key) = self.agdx.assemble(envelope, content_type)?;
buffer.messages.push(message);
let batch = std::mem::take(&mut buffer.messages);
buffer.first_at = None;
self.agdx
.laser
.send_batch(&self.agdx.topic, batch, Some(&partition_key))
.await?;
}
}
Ok(())
}
fn chunk(
&self,
body: Vec<u8>,
last: bool,
finish_reason: Option<String>,
usage: Option<TokenUsage>,
) -> AgentEnvelope {
let mut envelope = AgentEnvelope::chunk(
self.agdx.conversation,
self.agdx.source.clone(),
self.correlation,
self.channel,
self.sequence,
body,
);
if self.sequence == 0 {
envelope = envelope.with_operation(self.purpose.clone());
if let Some(deadline) = self.deadline_micros {
envelope = envelope.with_deadline_micros(deadline);
}
}
if let Some(target) = self.target.clone() {
envelope = envelope.with_target(target);
}
if last {
envelope.last = true;
envelope.finish_reason = finish_reason;
if let Some(usage) = usage {
envelope = envelope.with_usage(usage);
}
}
envelope
}
}
fn action_kind(kind: AgentKind) -> ActionKind {
match kind {
AgentKind::Command => ActionKind::Command,
AgentKind::Response => ActionKind::Response,
AgentKind::Event => ActionKind::Event,
AgentKind::Status => ActionKind::Status,
AgentKind::Error => ActionKind::Error,
_ => ActionKind::Send,
}
}
fn metadata_str<'e>(envelope: &'e AgentEnvelope, key: &str) -> Option<&'e str> {
match envelope.metadata.as_ref()?.get(key)? {
Value::Str(value) => Some(value.as_str()),
_ => None,
}
}
fn ensure_chunk_body_within_cap(body: &[u8]) -> Result<(), LaserError> {
if body.len() > MAX_CHUNK_BODY_BYTES {
return Err(LaserError::Invalid(format!(
"chunk body is {}B, exceeds cap {}B",
body.len(),
MAX_CHUNK_BODY_BYTES
)));
}
Ok(())
}
pub(crate) fn agdx_headers(
envelope: &AgentEnvelope,
content_type: ContentType,
) -> Result<BTreeMap<HeaderKey, HeaderValue>, LaserError> {
let mut headers = BTreeMap::new();
headers.insert(
HeaderKey::from_str(AGENT_VERSION)?,
HeaderValue::from(AGENT_OP_VERSION),
);
headers.insert(
HeaderKey::from_str(CONTENT_TYPE)?,
HeaderValue::from(content_type.code()),
);
headers.insert(
HeaderKey::from_str(CONVERSATION_ID)?,
HeaderValue::from(envelope.conversation.as_u128()),
);
if let Some(target) = &envelope.target {
headers.insert(
HeaderKey::from_str(TARGET_AGENT_ID)?,
HeaderValue::from_str(target.as_str())?,
);
}
Ok(headers)
}
#[cfg(test)]
mod tests {
use super::*;
use laser_wire::agent::OPERATION_CHAT;
use laser_wire::fixtures::assert_matches;
use serde::Serialize;
use serde::Serializer;
use std::env;
use std::fs;
use std::path::PathBuf;
#[test]
fn given_agdx_headers_when_stamped_then_should_pin_the_typed_encodings() {
let (record, conversation, source, correlation) = ids();
let envelope =
AgentEnvelope::command(record, conversation, source, correlation, b"x".to_vec())
.with_target("target-agent".parse().expect("valid agent id"));
let headers = agdx_headers(&envelope, ContentType::Json).expect("headers stamp");
let value = |key: &str| {
headers
.get(&HeaderKey::from_str(key).expect("key parses"))
.expect("header present")
};
assert_eq!(value(AGENT_VERSION).as_bytes(), [1, 0, 0, 0]);
assert_eq!(value(CONTENT_TYPE).as_bytes(), [1]);
assert_eq!(
value(CONVERSATION_ID).as_uint128().expect("uint128"),
conversation.as_u128()
);
assert_eq!(
value(CONVERSATION_ID).as_bytes(),
conversation.as_u128().to_le_bytes()
);
assert_eq!(
value(TARGET_AGENT_ID).as_str().expect("string"),
"target-agent"
);
assert_eq!(envelope.conversation.to_string().len(), 26);
}
#[test]
fn given_the_canonical_command_when_encoded_then_should_match_the_corpus() {
let (record, conversation, source, correlation) = ids();
let command = AgentEnvelope::command(
record,
conversation,
source,
correlation,
br#"{"ask":"plan the trip"}"#.to_vec(),
)
.with_target("target-agent".parse().expect("valid agent id"))
.with_idempotency_key("order-123-attempt-2".parse().expect("valid key"))
.with_deadline_micros(1_717_171_777_000_000)
.with_operation(OPERATION_CHAT)
.with_metadata("priority", "high");
let encoded = encode_named(&command).expect("encodes");
assert_matches("agent_command.bin", &encoded);
}
#[test]
fn given_an_oversized_chunk_body_when_checked_then_should_reject_before_publish() {
let body = vec![0u8; MAX_CHUNK_BODY_BYTES + 1];
let error = ensure_chunk_body_within_cap(&body).expect_err("oversized body rejects");
assert!(matches!(error, LaserError::Invalid(_)));
}
#[test]
fn given_the_canonical_record_when_encoded_then_should_match_record_fixture() {
let (record, conversation, source, correlation) = ids();
let envelope = AgentEnvelope::command(
record,
conversation,
source,
correlation,
br#"{"ask":"plan the trip"}"#.to_vec(),
)
.with_target("target-agent".parse().expect("valid agent id"))
.with_idempotency_key("order-123-attempt-2".parse().expect("valid key"))
.with_deadline_micros(1_717_171_777_000_000)
.with_operation(OPERATION_CHAT)
.with_metadata("priority", "high");
let headers = agdx_headers(&envelope, ContentType::Json).expect("headers stamp");
let payload = encode_named(&envelope).expect("payload encodes");
let record =
CanonicalAlpRecord::from_headers(envelope.conversation.to_string(), headers, payload);
let encoded = encode_named(&record).expect("record fixture encodes");
assert_record_fixture("agent_record.bin", &encoded);
}
#[derive(Serialize)]
struct CanonicalAlpRecord {
partition_key: String,
headers: BTreeMap<String, CanonicalHeader>,
payload: Binary,
}
impl CanonicalAlpRecord {
fn from_headers(
partition_key: String,
headers: BTreeMap<HeaderKey, HeaderValue>,
payload: Vec<u8>,
) -> Self {
let mut canonical_headers = BTreeMap::new();
for (key, value) in headers {
let key = key.as_str().expect("header key is utf8").to_owned();
let kind = match key.as_str() {
AGENT_VERSION => "u32",
CONTENT_TYPE => "u8",
CONVERSATION_ID => "uint128",
TARGET_AGENT_ID => "string",
_ => "bytes",
};
canonical_headers.insert(
key,
CanonicalHeader {
kind,
bytes: Binary(value.as_bytes().to_vec()),
},
);
}
Self {
partition_key,
headers: canonical_headers,
payload: Binary(payload),
}
}
}
#[derive(Serialize)]
struct CanonicalHeader {
kind: &'static str,
bytes: Binary,
}
struct Binary(Vec<u8>);
impl Serialize for Binary {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
serializer.serialize_bytes(&self.0)
}
}
fn assert_record_fixture(name: &str, encoded: &[u8]) {
let path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.parent()
.expect("sdk crate has workspace parent")
.join("wire")
.join("fixtures")
.join(name);
if env::var("AGDX_WIRE_FIXTURES_REGEN").is_ok() {
fs::write(&path, encoded).expect("write record fixture");
}
let golden = fs::read(&path).unwrap_or_else(|error| {
panic!("read fixture {name}: {error} (regen with AGDX_WIRE_FIXTURES_REGEN=1)")
});
assert_eq!(
encoded, golden,
"fixture `{name}` drifted from the canonical AGDX record"
);
}
fn ids() -> (RecordId, ConversationId, AgentId, CorrelationId) {
(
RecordId::from_u128(0x0190_3c1f_aa00_0000_0000_0000_0000_0001),
ConversationId::from_u128(0x0190_3c1f_aa00_0000_0000_0000_0000_0002),
"source-agent".parse().expect("valid agent id"),
CorrelationId::from_u128(0x0190_3c1f_aa00_0000_0000_0000_0000_0005),
)
}
}