laser_sdk/agent/
agdx.rs

1use crate::error::LaserError;
2use crate::govern::{ActionCounters, ActionKind, GovernedAction};
3use crate::laser::Laser;
4use crate::provenance::AgentTopic;
5use crate::types::MintUlid;
6use iggy::prelude::{HeaderKey, HeaderValue};
7use laser_wire::agent::{
8    AgentEnvelope, AgentErrorBody, AgentId, AgentKind, ChannelId, ConversationId, CorrelationId,
9    IdempotencyKey, LogPosition, RecordId, TaskState, TokenUsage, validate,
10};
11#[cfg(feature = "sign")]
12use laser_wire::codes::AGENT_OP_VERSION;
13use laser_wire::content::ContentType;
14use laser_wire::framing::{decode_named, encode_named};
15use laser_wire::query::Value;
16use std::collections::BTreeMap;
17use std::time::Duration;
18
19// Producer-side chunking guidance, DRAFT until the benchmark step pins them:
20// flush a chunk at the byte target or the linger bound, whichever first, and
21// never exceed the hard chunk cap.
22/// Draft chunk flush target, in body bytes.
23pub const DEFAULT_CHUNK_FLUSH_BYTES: usize = 512;
24/// Draft chunk linger bound, in milliseconds.
25pub const DEFAULT_CHUNK_LINGER_MS: u64 = 20;
26/// Draft hard cap on one chunk's body.
27pub const MAX_CHUNK_BODY_BYTES: usize = 64 * 1024;
28
29/// The typed AGDX producer over one agent topic: every send is a validated
30/// [`AgentEnvelope`] (invalid envelopes are unrepresentable or rejected at
31/// publish time), encoded as named-field CBOR, stamped with the routing
32/// headers (`agdx.av` u32, `agdx.ct` u8, the conversation as a Crockford string,
33/// and the target as the agent's name string), and partition-keyed by the
34/// conversation's canonical base32 form so one conversation stays ordered.
35#[derive(Clone)]
36pub struct Agdx {
37    laser: Laser,
38    topic: String,
39    source: AgentId,
40    conversation: ConversationId,
41    lane_guard: Option<std::sync::Arc<crate::agent::session::LaneGuard>>,
42    stream_generation: Option<u64>,
43}
44
45impl Laser {
46    /// A typed AGDX producer publishing as `source` within `conversation` on
47    /// `topic`. The topic is resolved to its name at construction, so a runtime
48    /// `AgentTopic::Custom` built from a borrowed identifier is accepted.
49    pub fn agdx(
50        &self,
51        topic: AgentTopic<'_>,
52        source: impl Into<AgentId>,
53        conversation: ConversationId,
54    ) -> Agdx {
55        let source = source.into();
56        Agdx {
57            laser: self.clone(),
58            topic: topic.topic_string(),
59            source,
60            conversation,
61            lane_guard: None,
62            stream_generation: None,
63        }
64    }
65}
66
67impl Agdx {
68    pub(crate) fn with_lane_guard(
69        mut self,
70        guard: std::sync::Arc<crate::agent::session::LaneGuard>,
71    ) -> Self {
72        self.lane_guard = Some(guard);
73        self
74    }
75
76    pub(crate) fn with_stream_generation(mut self, generation: u64) -> Self {
77        self.stream_generation = Some(generation);
78        self
79    }
80
81    /// A `command`: expects a reply or effect under `correlation`.
82    pub fn command(&self, correlation: CorrelationId, body: Vec<u8>) -> AgdxSend<'_> {
83        self.send_of(AgentEnvelope::command(
84            RecordId::mint(),
85            self.conversation,
86            self.source.clone(),
87            correlation,
88            body,
89        ))
90    }
91
92    /// A `response`: the paired answer to a command (same `correlation`).
93    pub fn respond(&self, correlation: CorrelationId, body: Vec<u8>) -> AgdxSend<'_> {
94        self.send_of(AgentEnvelope::response(
95            RecordId::mint(),
96            self.conversation,
97            self.source.clone(),
98            correlation,
99            body,
100        ))
101    }
102
103    /// An `event`: expects nothing.
104    pub fn emit(&self, body: Vec<u8>) -> AgdxSend<'_> {
105        self.send_of(AgentEnvelope::event(
106            RecordId::mint(),
107            self.conversation,
108            self.source.clone(),
109            body,
110        ))
111    }
112
113    /// A `status` signal discriminated by `operation` (`task` | `card` |
114    /// `progress`). Task updates additionally chain
115    /// [`with_correlation`](AgdxSend::with_correlation) and
116    /// [`with_task_state`](AgdxSend::with_task_state).
117    pub fn status(&self, operation: impl Into<String>) -> AgdxSend<'_> {
118        self.send_of(AgentEnvelope::status(
119            RecordId::mint(),
120            self.conversation,
121            self.source.clone(),
122            operation,
123        ))
124    }
125
126    /// An `error` terminal for `correlation`. The body is the encoded
127    /// [`AgentErrorBody`] (so `agdx.ct` is forced to cbor).
128    pub fn fail(
129        &self,
130        correlation: CorrelationId,
131        error: &AgentErrorBody,
132    ) -> Result<AgdxSend<'_>, LaserError> {
133        let body = encode_named(error)?;
134        Ok(self
135            .send_of(AgentEnvelope::error(
136                RecordId::mint(),
137                self.conversation,
138                self.source.clone(),
139                correlation,
140                body,
141            ))
142            .content_type(ContentType::Cbor))
143    }
144
145    /// A chunk-stream writer under `correlation` on a fresh channel. The
146    /// `purpose` is the pinned chunk-stream vocabulary (`chat` | `reasoning`
147    /// | `tool_args`), declared on the opening chunk.
148    pub fn stream(&self, correlation: CorrelationId, purpose: impl Into<String>) -> AgdxStream {
149        AgdxStream {
150            agdx: self.clone(),
151            correlation,
152            channel: ChannelId::mint(),
153            purpose: purpose.into(),
154            sequence: 0,
155            deadline_micros: None,
156            target: None,
157            content_type: ContentType::Raw,
158            buffer: None,
159        }
160    }
161
162    /// Human-in-the-loop interrupt/resume. Pauses on a human: publishes a prompt
163    /// `command` under a fresh interrupt correlation on this producer's topic,
164    /// then awaits the human's correlated `response` on `reply_topic` up to
165    /// `timeout` and returns its body. A responder answers with
166    /// [`AgentCtx::respond_input`](crate::agent::AgentCtx::respond_input) (a
167    /// `response`) or rejects with an `error`, which surfaces here as
168    /// [`LaserError::Rejected`]. Composes existing verbs, so it adds nothing to
169    /// the wire. It blocks the caller until the response lands or the timeout
170    /// elapses, which is the point: the task is genuinely paused on a human.
171    /// The prompt is addressed to every agent (`agdx.to = *`), so on a shared
172    /// session topic every listening agent receives it. Use
173    /// [`request_input_from`](Self::request_input_from) to ask one agent.
174    pub async fn request_input(
175        &self,
176        reply_topic: AgentTopic<'_>,
177        prompt: impl Into<Vec<u8>>,
178        timeout: Duration,
179    ) -> Result<Vec<u8>, LaserError> {
180        self.request_input_addressed(None, reply_topic, prompt.into(), timeout)
181            .await
182    }
183
184    /// [`request_input`](Self::request_input) with the prompt addressed to
185    /// `target`, so only that agent answers it on a shared session topic.
186    pub async fn request_input_from(
187        &self,
188        target: impl Into<AgentId>,
189        reply_topic: AgentTopic<'_>,
190        prompt: impl Into<Vec<u8>>,
191        timeout: Duration,
192    ) -> Result<Vec<u8>, LaserError> {
193        let target = target.into();
194        self.request_input_addressed(Some(target), reply_topic, prompt.into(), timeout)
195            .await
196    }
197
198    async fn request_input_addressed(
199        &self,
200        target: Option<AgentId>,
201        reply_topic: AgentTopic<'_>,
202        prompt: Vec<u8>,
203        timeout: Duration,
204    ) -> Result<Vec<u8>, LaserError> {
205        let interrupt = CorrelationId::mint();
206        // Seed the reply reader at the topic tail before sending the prompt, so it
207        // reads only the human's response rather than the topic's history.
208        let mut reader = self
209            .laser
210            .agdx_reply_reader(reply_topic, Some(&self.source))
211            .await?;
212        let command = self.command(interrupt, prompt);
213        match target {
214            Some(target) => command.with_target(target).send().await?,
215            None => command.send().await?,
216        };
217        let reply = self
218            .laser
219            .await_agdx_reply(&mut reader, interrupt, timeout)
220            .await?;
221        if reply.kind == AgentKind::Error {
222            let message = decode_named::<AgentErrorBody>(&reply.body)
223                .ok()
224                .and_then(|body| body.message)
225                .unwrap_or_else(|| "the input request was rejected".to_owned());
226            return Err(LaserError::Rejected(message));
227        }
228        Ok(reply.body)
229    }
230
231    // Publish an envelope built elsewhere, unchanged, on this producer's topic.
232    pub(crate) async fn publish_envelope(
233        &self,
234        envelope: AgentEnvelope,
235    ) -> Result<AgdxReceipt, LaserError> {
236        self.publish(envelope, ContentType::Raw).await
237    }
238
239    fn send_of(&self, envelope: AgentEnvelope) -> AgdxSend<'_> {
240        AgdxSend {
241            agdx: self,
242            envelope,
243            content_type: ContentType::Raw,
244            #[cfg(feature = "sign")]
245            sign_key: None,
246            claim_check: None,
247        }
248    }
249
250    #[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"))]
251    async fn publish(
252        &self,
253        envelope: AgentEnvelope,
254        content_type: ContentType,
255    ) -> Result<AgdxReceipt, LaserError> {
256        let record = envelope.record;
257        // A declared per-agent topic layout moves addressed work off the lane.
258        // The lane identity check covers lane records only.
259        let destination = self.laser.agent_destination(
260            &self.topic,
261            Some(envelope.kind),
262            envelope.target.as_ref(),
263        );
264        let (message, partitioning) = self.assemble(envelope, content_type)?;
265        let sent = self
266            .laser
267            .send_batch_partitioned_on_checked(
268                self.laser.stream_required()?,
269                destination.as_deref().unwrap_or(&self.topic),
270                vec![message],
271                partitioning.clone(),
272                || async {
273                    if destination.is_some() {
274                        return Ok(None);
275                    }
276                    let (stream_id, topic_id, partitions) = if let Some(guard) = &self.lane_guard {
277                        let identity = guard.check(&self.laser, self.conversation.into()).await?;
278                        (identity.stream_id, identity.topic_id, identity.partitions)
279                    } else if let Some(generation) = self.stream_generation {
280                        use iggy::prelude::StreamClient;
281                        let details = self
282                            .laser
283                            .client()
284                            .get_stream(&iggy::prelude::Identifier::named(
285                                self.laser.stream_required()?,
286                            )?)
287                            .await?
288                            .ok_or_else(|| {
289                                LaserError::Session(laser_wire::session::SessionError::Stale(
290                                    "the heartbeat stream does not exist".to_owned(),
291                                ))
292                            })?;
293                        if details.created_at.as_micros() != generation {
294                            return Err(LaserError::Session(
295                                laser_wire::session::SessionError::Stale(
296                                    "the heartbeat stream generation changed".to_owned(),
297                                ),
298                            ));
299                        }
300                        let topic = details
301                            .topics
302                            .iter()
303                            .find(|topic| topic.name == self.topic)
304                            .filter(|topic| topic.partitions_count > 0)
305                            .ok_or_else(|| {
306                                LaserError::Session(laser_wire::session::SessionError::Stale(
307                                    "the heartbeat topic does not exist".to_owned(),
308                                ))
309                            })?;
310                        (details.id, topic.id, topic.partitions_count)
311                    } else {
312                        return Ok(None);
313                    };
314                    let route = if partitioning.kind == iggy_common::PartitioningKind::MessagesKey {
315                        iggy::prelude::Partitioning::partition_id(
316                            iggy_common::calculate_32(&partitioning.value) % partitions,
317                        )
318                    } else {
319                        partitioning.clone()
320                    };
321                    Ok(Some((stream_id, topic_id, route)))
322                },
323            )
324            .await?;
325        let confirmation = sent.confirmations.first();
326        Ok(AgdxReceipt {
327            record,
328            partition_id: confirmation.map(|confirmation| confirmation.partition_id),
329            offset: confirmation.map(|confirmation| confirmation.base_offset),
330        })
331    }
332
333    // Validate and encode one envelope into a ready message plus its
334    // partitioning: the single assembly the per-send publish and the
335    // buffered stream's batch append share, so buffering changes when bytes
336    // move, never what they are.
337    fn assemble(
338        &self,
339        envelope: AgentEnvelope,
340        content_type: ContentType,
341    ) -> Result<(iggy::prelude::IggyMessage, iggy::prelude::Partitioning), LaserError> {
342        validate(&envelope)?;
343        let payload = encode_named(&envelope)?;
344        // The session topics carry an addressee on every record, `*` when
345        // untargeted, so the addressee filter of a role group keeps them.
346        let headers = agdx_headers(
347            &envelope,
348            content_type,
349            self.topic == laser_wire::topics::AGENT_SESSIONS
350                || self.topic == laser_wire::topics::AGENT_CONTROL,
351        )?;
352        let partitioning = self.partitioning(envelope.kind, envelope.target.as_ref())?;
353        let message = iggy::prelude::IggyMessage::builder()
354            .payload(bytes::Bytes::from(payload))
355            .user_headers(headers)
356            .build()?;
357        Ok((message, partitioning))
358    }
359
360    // The partition a record of `kind` addressed to `target` lands on under
361    // this stream's declared layout.
362    fn partitioning(
363        &self,
364        kind: laser_wire::agent::AgentKind,
365        target: Option<&AgentId>,
366    ) -> Result<iggy::prelude::Partitioning, LaserError> {
367        let layout = self
368            .laser
369            .default_stream()
370            .and_then(|stream| self.laser.layout(stream));
371        crate::agent::partitioning::AgentPartitioning::resolve(
372            layout.as_ref(),
373            &self.topic,
374            kind,
375            target,
376        )
377        .into_partitioning(self.conversation)
378    }
379}
380
381/// Where one published envelope was committed.
382#[derive(Debug, Clone, Copy, PartialEq, Eq)]
383pub struct AgdxReceipt {
384    /// The envelope's record id. Chunks have none.
385    pub record: Option<RecordId>,
386    /// The partition the record landed on, when the server confirmed it.
387    pub partition_id: Option<u32>,
388    /// The record's offset, when the server confirmed it.
389    pub offset: Option<u64>,
390}
391
392/// One pending AGDX send: the envelope built by its verb, refined by the
393/// `with_*` setters, validated and published by [`send`](Self::send).
394#[must_use = "an unsent envelope does nothing until you call .send().await"]
395pub struct AgdxSend<'a> {
396    agdx: &'a Agdx,
397    envelope: AgentEnvelope,
398    content_type: ContentType,
399    /// When set, the envelope is signed with this key just before encoding (the
400    /// last step, so all the `with_*` refinements are covered). A verifying
401    /// consumer rejects an unsigned or unverified record on a control or effect
402    /// topic (the mandatory-verification gate), which is what makes a cancel or
403    /// quarantine authorizable.
404    #[cfg(feature = "sign")]
405    sign_key: Option<&'a crate::sign::SigningKey>,
406    /// When set, a body at or over the threshold is externalized to the store
407    /// and replaced by the `BodyRef` capsule at send (content-type `ref`).
408    claim_check: Option<(&'a dyn crate::blob::BlobStore, usize)>,
409}
410
411impl<'a> AgdxSend<'a> {
412    /// Use `record` as the envelope's record id, so a retried send repeats the
413    /// same record.
414    pub(crate) fn with_record(mut self, record: RecordId) -> Self {
415        self.envelope.record = Some(record);
416        self
417    }
418
419    /// Mark the envelope as part of a child session whose parent is `parent`
420    /// and whose tree is rooted at `root`.
421    pub fn with_ancestry(
422        mut self,
423        parent: Option<laser_wire::agent::ConversationId>,
424        root: Option<laser_wire::agent::ConversationId>,
425    ) -> Self {
426        self.envelope.parent = parent;
427        self.envelope.root = root;
428        self
429    }
430
431    /// Narrow delivery to one agent within the shared topic (routing, never an ACL).
432    pub fn with_target(mut self, target: impl Into<AgentId>) -> Self {
433        self.envelope = self.envelope.with_target(target.into());
434        self
435    }
436
437    /// Stamp the causal parent (identity, plus the locator when known).
438    pub fn with_cause(mut self, cause: RecordId, cause_at: Option<LogPosition>) -> Self {
439        self.envelope = self.envelope.with_cause(cause, cause_at);
440        self
441    }
442
443    /// Pair with a correlation id (required by status task updates).
444    pub fn with_correlation(mut self, correlation: CorrelationId) -> Self {
445        self.envelope = self.envelope.with_correlation(correlation);
446        self
447    }
448
449    /// Attach a business idempotency key.
450    pub fn with_idempotency_key(mut self, key: IdempotencyKey) -> Self {
451        self.envelope = self.envelope.with_idempotency_key(key);
452        self
453    }
454
455    /// Declare the drop-dead time, epoch micros.
456    pub fn with_deadline_micros(mut self, deadline_micros: u64) -> Self {
457        self.envelope = self.envelope.with_deadline_micros(deadline_micros);
458        self
459    }
460
461    /// Attach a task state (status task updates require it).
462    pub fn with_task_state(mut self, state: TaskState) -> Self {
463        self.envelope = self.envelope.with_task_state(state);
464        self
465    }
466
467    /// Set the OTel operation name (open vocabulary kinds only).
468    pub fn with_operation(mut self, operation: impl Into<String>) -> Self {
469        self.envelope = self.envelope.with_operation(operation);
470        self
471    }
472
473    /// Set why the model stopped generating.
474    pub fn with_finish_reason(mut self, reason: impl Into<String>) -> Self {
475        self.envelope.finish_reason = Some(reason.into());
476        self
477    }
478
479    /// Set the OTel tool name.
480    pub fn with_tool(mut self, tool: impl Into<String>) -> Self {
481        self.envelope = self.envelope.with_tool(tool);
482        self
483    }
484
485    /// Attach token accounting (advisory).
486    pub fn with_usage(mut self, usage: TokenUsage) -> Self {
487        self.envelope = self.envelope.with_usage(usage);
488        self
489    }
490
491    /// Add one AGDX-native metadata entry.
492    pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<Value>) -> Self {
493        self.envelope = self.envelope.with_metadata(key, value);
494        self
495    }
496
497    /// Mark a status terminal (`last = true` on the final task update).
498    pub fn last(mut self) -> Self {
499        self.envelope.last = true;
500        self
501    }
502
503    /// Declare the body's codec (`agdx.ct`). Defaults to raw. An `error`
504    /// envelope always carries a CBOR [`AgentErrorBody`], so any other content
505    /// type on [`Agdx::fail`] fails at send.
506    pub fn content_type(mut self, content_type: ContentType) -> Self {
507        self.content_type = content_type;
508        self
509    }
510
511    /// Attach the body for the kinds whose verb does not take one (`status`,
512    /// notably the `card` body of a registry advertisement). The body-carrying
513    /// verbs (`command`/`respond`/`emit`/`fail`) set it directly, so this is for
514    /// refining a `status`.
515    pub fn body(mut self, body: impl Into<Vec<u8>>) -> Self {
516        self.envelope.body = body.into();
517        self
518    }
519
520    /// Sign the envelope with `key` just before sending (after every `with_*`
521    /// refinement, so the signature covers the final shape). The consuming SDK
522    /// verifies it against the enrolled key registry, the authorship and
523    /// authorization gate. Requires the `sign` feature.
524    #[cfg(feature = "sign")]
525    pub fn signed_by(mut self, key: &'a crate::sign::SigningKey) -> Self {
526        self.sign_key = Some(key);
527        self
528    }
529
530    /// Claim-check the body against `store` when it is at or over
531    /// `threshold_bytes` at send time: the body is externalized and replaced
532    /// by the [`BodyRef`](laser_wire::agent::BodyRef) capsule, content-type
533    /// `ref`. Applied before signing, so a signature covers the capsule the
534    /// log actually carries. Readers resolve with
535    /// [`AgentMessage::resolve_body`](crate::agent::AgentMessage::resolve_body).
536    pub fn claim_check(
537        mut self,
538        store: &'a dyn crate::blob::BlobStore,
539        threshold_bytes: usize,
540    ) -> Self {
541        self.claim_check = Some((store, threshold_bytes));
542        self
543    }
544
545    /// Validate, encode, stamp the headers, and publish. Returns the minted
546    /// record id (chunks have none). When [`signed_by`](Self::signed_by) set a
547    /// key, the envelope is signed here, last, so the signature covers every
548    /// refinement.
549    pub async fn send(self) -> Result<Option<RecordId>, LaserError> {
550        Ok(self.send_receipt().await?.record)
551    }
552
553    /// Like [`send`](Self::send), and also return where the record was
554    /// committed.
555    pub async fn send_receipt(self) -> Result<AgdxReceipt, LaserError> {
556        let mut envelope = self.envelope;
557        let mut content_type = self.content_type;
558        if envelope.kind == AgentKind::Error && content_type != ContentType::Cbor {
559            return Err(LaserError::Invalid(
560                "an error envelope carries a CBOR AgentErrorBody, so its content type cannot change"
561                    .to_owned(),
562            ));
563        }
564        // The pre-effect policy hook, before claim-check and signing so a
565        // modified body is what gets externalized and what the signature covers.
566        #[cfg(feature = "sign")]
567        let will_sign = self.sign_key.is_some();
568        #[cfg(not(feature = "sign"))]
569        let will_sign = false;
570        let correlation = envelope.correlation.map(|value| value.to_string());
571        let action = GovernedAction {
572            kind: action_kind(envelope.kind),
573            stream: self.agdx.laser.stream_required()?,
574            topic: &self.agdx.topic,
575            source: Some(envelope.source.as_str()),
576            target: envelope.target.as_ref().map(|target| target.as_str()),
577            conversation: Some(envelope.conversation.into()),
578            correlation: correlation.as_deref(),
579            operation: envelope.operation.as_deref(),
580            tool: envelope.tool.as_deref(),
581            on_behalf_of: metadata_str(&envelope, laser_wire::agent::METADATA_DELEGATED_BY),
582            purpose: metadata_str(&envelope, laser_wire::agent::METADATA_PURPOSE),
583            data_classification: metadata_str(
584                &envelope,
585                laser_wire::agent::METADATA_DATA_CLASSIFICATION,
586            ),
587            payload: &envelope.body,
588            signed: will_sign,
589            counters: ActionCounters::default(),
590        };
591        if let Some(modified) = self.agdx.laser.govern(action).await? {
592            envelope.body = modified;
593        }
594        if let Some((store, threshold)) = self.claim_check {
595            let body = std::mem::take(&mut envelope.body);
596            let (body, replaced) = crate::blob::check_in(store, threshold, body).await?;
597            envelope.body = body;
598            if let Some(replaced) = replaced {
599                content_type = replaced;
600            }
601        }
602        #[cfg(feature = "sign")]
603        let envelope = if let Some(key) = self.sign_key {
604            // Bind the interpretation attributes the log will stamp out of band
605            // (`agdx.ct` = the content-type code, `agdx.av` = the wire version)
606            // into the signature, so neither can be flipped on the signed record.
607            let context = laser_wire::agent::SignatureContext {
608                content_type: Some(content_type.code()),
609                agent_version: Some(AGENT_OP_VERSION),
610            };
611            let signature = key.sign_with_context(&envelope, context)?;
612            envelope.with_signature(signature)
613        } else {
614            envelope
615        };
616        self.agdx.publish(envelope, content_type).await
617    }
618}
619
620/// A chunk-stream writer: auto-incremented `sequence`, the purpose and the
621/// abandonment bound on the opening chunk, one terminal (`finish` or `fail`).
622/// Dropping the writer without a terminal is the producer-death case readers
623/// abandon by deadline.
624pub struct AgdxStream {
625    agdx: Agdx,
626    correlation: CorrelationId,
627    channel: ChannelId,
628    purpose: String,
629    sequence: u64,
630    deadline_micros: Option<u64>,
631    target: Option<AgentId>,
632    content_type: ContentType,
633    buffer: Option<ChunkBuffer>,
634}
635
636// The opt-in write buffer: assembled chunk messages held until `max_chunks`
637// or `linger` trips, then appended as ONE batch. Sequences are assigned at
638// write, so reassembly semantics are untouched.
639struct ChunkBuffer {
640    max_chunks: usize,
641    linger: std::time::Duration,
642    first_at: Option<std::time::Instant>,
643    messages: Vec<iggy::prelude::IggyMessage>,
644}
645
646impl AgdxStream {
647    /// The stream's channel id.
648    pub fn channel(&self) -> ChannelId {
649        self.channel
650    }
651
652    /// Declare the reader-local abandonment bound (rides the opening chunk,
653    /// so set it before the first write).
654    pub fn with_deadline_micros(mut self, deadline_micros: u64) -> Self {
655        self.deadline_micros = deadline_micros.into();
656        self
657    }
658
659    /// Narrow delivery to one agent within the shared topic.
660    pub fn with_target(mut self, target: AgentId) -> Self {
661        self.target = Some(target);
662        self
663    }
664
665    /// Declare the chunk bodies' codec (`agdx.ct`). Defaults to raw.
666    pub fn content_type(mut self, content_type: ContentType) -> Self {
667        self.content_type = content_type;
668        self
669    }
670
671    /// Buffer writes and append them as ONE batch when `max_chunks` fill or
672    /// `linger` has passed since the first buffered chunk, whichever first.
673    /// Opt-in for chunk-heavy token streams. The unbuffered stream publishes
674    /// each write exactly as before. The linger is checked at write (no
675    /// background task drives this writer), so a stalled producer holds its
676    /// buffered chunks until the next `write`, an explicit
677    /// [`flush`](Self::flush), or the terminal, which always flushes.
678    #[must_use]
679    pub fn buffered(mut self, max_chunks: usize, linger: std::time::Duration) -> Self {
680        self.buffer = Some(ChunkBuffer {
681            max_chunks: max_chunks.max(1),
682            linger,
683            first_at: None,
684            messages: Vec::new(),
685        });
686        self
687    }
688
689    /// Publish the next chunk. The opening chunk (`sequence` 0) carries the
690    /// stream purpose and the abandonment bound.
691    pub async fn write(&mut self, body: Vec<u8>) -> Result<(), LaserError> {
692        ensure_chunk_body_within_cap(&body)?;
693        let envelope = self.chunk(body, false, None, None);
694        self.sequence += 1;
695        match &mut self.buffer {
696            None => {
697                self.agdx.publish(envelope, self.content_type).await?;
698            }
699            Some(buffer) => {
700                let (message, _) = self.agdx.assemble(envelope, self.content_type)?;
701                buffer.first_at.get_or_insert_with(std::time::Instant::now);
702                buffer.messages.push(message);
703                let linger_tripped = buffer
704                    .first_at
705                    .is_some_and(|first| first.elapsed() >= buffer.linger);
706                if buffer.messages.len() >= buffer.max_chunks || linger_tripped {
707                    self.flush().await?;
708                }
709            }
710        }
711        Ok(())
712    }
713
714    /// Append everything buffered as one batch. A no-op unbuffered or empty.
715    pub async fn flush(&mut self) -> Result<(), LaserError> {
716        let Some(buffer) = &mut self.buffer else {
717            return Ok(());
718        };
719        if buffer.messages.is_empty() {
720            return Ok(());
721        }
722        let batch = std::mem::take(&mut buffer.messages);
723        buffer.first_at = None;
724        let partitioning = self
725            .agdx
726            .partitioning(laser_wire::agent::AgentKind::Chunk, self.target.as_ref())?;
727        let destination = self.destination();
728        self.agdx
729            .laser
730            .send_batch_partitioned_on(
731                self.agdx.laser.stream_required()?,
732                &destination,
733                batch,
734                partitioning,
735            )
736            .await
737            .map(|_| ())
738    }
739
740    /// Publish the terminal chunk (`last = true`) with the reason the stream
741    /// ended and the whole-stream accounting. A buffered stream appends its
742    /// held chunks and the terminal together, one batch.
743    pub async fn finish(
744        mut self,
745        finish_reason: impl Into<String>,
746        usage: Option<TokenUsage>,
747    ) -> Result<(), LaserError> {
748        let envelope = self.chunk(Vec::new(), true, Some(finish_reason.into()), usage);
749        self.send_terminal(envelope, self.content_type).await
750    }
751
752    /// Terminate the stream with a `kind = error` terminal carrying this
753    /// channel and the next sequence.
754    pub async fn fail(mut self, error: &AgentErrorBody) -> Result<(), LaserError> {
755        let body = encode_named(error)?;
756        let mut envelope = AgentEnvelope::error(
757            RecordId::mint(),
758            self.agdx.conversation,
759            self.agdx.source.clone(),
760            self.correlation,
761            body,
762        );
763        envelope.channel = Some(self.channel);
764        envelope.sequence = Some(self.sequence);
765        if let Some(target) = self.target.clone() {
766            envelope = envelope.with_target(target);
767        }
768        self.send_terminal(envelope, ContentType::Cbor).await
769    }
770
771    // The one terminal path: a buffered stream appends its held chunks and
772    // the terminal as ONE batch (ordered, sequence-stamped at write), an
773    // unbuffered one publishes the terminal alone.
774    async fn send_terminal(
775        &mut self,
776        envelope: AgentEnvelope,
777        content_type: ContentType,
778    ) -> Result<(), LaserError> {
779        match &mut self.buffer {
780            None => {
781                self.agdx.publish(envelope, content_type).await?;
782            }
783            Some(buffer) => {
784                let (message, partitioning) = self.agdx.assemble(envelope, content_type)?;
785                buffer.messages.push(message);
786                let batch = std::mem::take(&mut buffer.messages);
787                buffer.first_at = None;
788                let destination = self.destination();
789                self.agdx
790                    .laser
791                    .send_batch_partitioned_on(
792                        self.agdx.laser.stream_required()?,
793                        &destination,
794                        batch,
795                        partitioning,
796                    )
797                    .await?;
798            }
799        }
800        Ok(())
801    }
802
803    // Where this stream's buffered chunks land: the addressee's declared
804    // topic under a per-agent topic layout, else the producer's topic.
805    fn destination(&self) -> String {
806        self.agdx
807            .laser
808            .agent_destination(
809                &self.agdx.topic,
810                Some(laser_wire::agent::AgentKind::Chunk),
811                self.target.as_ref(),
812            )
813            .unwrap_or_else(|| self.agdx.topic.clone())
814    }
815
816    fn chunk(
817        &self,
818        body: Vec<u8>,
819        last: bool,
820        finish_reason: Option<String>,
821        usage: Option<TokenUsage>,
822    ) -> AgentEnvelope {
823        let mut envelope = AgentEnvelope::chunk(
824            self.agdx.conversation,
825            self.agdx.source.clone(),
826            self.correlation,
827            self.channel,
828            self.sequence,
829            body,
830        );
831        if self.sequence == 0 {
832            envelope = envelope.with_operation(self.purpose.clone());
833            if let Some(deadline) = self.deadline_micros {
834                envelope = envelope.with_deadline_micros(deadline);
835            }
836        }
837        if let Some(target) = self.target.clone() {
838            envelope = envelope.with_target(target);
839        }
840        if last {
841            envelope.last = true;
842            envelope.finish_reason = finish_reason;
843            if let Some(usage) = usage {
844                envelope = envelope.with_usage(usage);
845            }
846        }
847        envelope
848    }
849}
850
851// The governed-action kind of an envelope kind. Chunk streams bypass
852// `AgdxSend` (their writer publishes directly), so the wildcard is unreachable
853// today and maps any future kind to a plain send rather than skipping the hook.
854fn action_kind(kind: AgentKind) -> ActionKind {
855    match kind {
856        AgentKind::Command => ActionKind::Command,
857        AgentKind::Response => ActionKind::Response,
858        AgentKind::Event => ActionKind::Event,
859        AgentKind::Status => ActionKind::Status,
860        AgentKind::Error => ActionKind::Error,
861        _ => ActionKind::Send,
862    }
863}
864
865// A string metadata entry of `envelope`, for the advisory policy inputs.
866fn metadata_str<'e>(envelope: &'e AgentEnvelope, key: &str) -> Option<&'e str> {
867    match envelope.metadata.as_ref()?.get(key)? {
868        Value::Str(value) => Some(value.as_str()),
869        _ => None,
870    }
871}
872
873fn ensure_chunk_body_within_cap(body: &[u8]) -> Result<(), LaserError> {
874    if body.len() > MAX_CHUNK_BODY_BYTES {
875        return Err(LaserError::Invalid(format!(
876            "chunk body is {}B, exceeds cap {}B",
877            body.len(),
878            MAX_CHUNK_BODY_BYTES
879        )));
880    }
881    Ok(())
882}
883
884// The AGDX routing headers from the shared wire encoder: `agdx.av` selects the
885// decoder before any body byte is read, `agdx.ct` names the inner body codec,
886// the conversation and its ancestry ride as canonical strings, and the author
887// and addressee ride as agent names, so projections and plain consumers route
888// without decoding the CBOR body.
889pub(crate) fn agdx_headers(
890    envelope: &AgentEnvelope,
891    content_type: ContentType,
892    broadcast: bool,
893) -> Result<BTreeMap<HeaderKey, HeaderValue>, LaserError> {
894    let headers =
895        laser_wire::headers::RecordHeaders::for_envelope(envelope, content_type, broadcast);
896    Ok(crate::provenance::header_map(&headers)?)
897}
898
899#[cfg(test)]
900mod tests {
901    use super::*;
902    use laser_wire::agent::OPERATION_CHAT;
903    use laser_wire::fixtures::assert_matches;
904    use laser_wire::headers::{
905        AGENT_ID, AGENT_VERSION, CONTENT_TYPE, CONVERSATION_ID, PARENT_CONVERSATION_ID,
906        ROOT_CONVERSATION_ID, TARGET_AGENT_ID,
907    };
908    use serde::Serialize;
909    use serde::Serializer;
910    use std::env;
911    use std::fs;
912    use std::path::PathBuf;
913    use std::str::FromStr;
914
915    #[tokio::test]
916    async fn given_a_fail_with_another_content_type_when_sent_then_should_refuse_it() {
917        let laser = Laser::from_client(iggy::prelude::IggyClient::default())
918            .with_default_stream("agdx-fail");
919        let agdx = laser.agdx(
920            AgentTopic::Sessions,
921            AgentId::from_str("worker").expect("a valid agent id"),
922            ConversationId::mint(),
923        );
924        let error = AgentErrorBody {
925            code: laser_wire::agent::AgentErrorCode::Internal,
926            message: Some("boom".to_owned()),
927            retryable: false,
928            detail: None,
929        };
930        let result = agdx
931            .fail(CorrelationId::mint(), &error)
932            .expect("the error body encodes")
933            .content_type(ContentType::Json)
934            .send_receipt()
935            .await;
936        assert!(matches!(result, Err(LaserError::Invalid(_))));
937    }
938
939    // Record conformance, the header half: the typed header encodings the
940    // runtime stamps, asserted byte-for-byte. The payload half is pinned
941    // against the shared corpus below.
942    #[test]
943    fn given_agdx_headers_when_stamped_then_should_pin_the_typed_encodings() {
944        let (record, conversation, source, correlation) = ids();
945        let envelope =
946            AgentEnvelope::command(record, conversation, source, correlation, b"x".to_vec())
947                .with_target("target-agent".parse().expect("valid agent id"));
948        let headers = agdx_headers(&envelope, ContentType::Json, false).expect("headers stamp");
949
950        let value = |key: &str| {
951            headers
952                .get(&HeaderKey::from_str(key).expect("key parses"))
953                .expect("header present")
954        };
955        // agdx.av: u32, little-endian.
956        assert_eq!(value(AGENT_VERSION).as_bytes(), [1, 0, 0, 0]);
957        // agdx.ct: one byte, json = 1.
958        assert_eq!(value(CONTENT_TYPE).as_bytes(), [1]);
959        assert_eq!(
960            value(CONVERSATION_ID).as_str().expect("string"),
961            conversation.to_string()
962        );
963        // The agent routing id is a string (the agent's name), not a numeric id.
964        assert_eq!(
965            value(TARGET_AGENT_ID).as_str().expect("string"),
966            "target-agent"
967        );
968        // The author rides beside the addressee so a header filter selects both.
969        assert_eq!(
970            value(AGENT_ID).as_str().expect("string"),
971            envelope.source.as_str()
972        );
973        // The partition key is the conversation's canonical base32 form.
974        assert_eq!(envelope.conversation.to_string().len(), 26);
975    }
976
977    #[test]
978    fn given_child_session_envelope_when_stamped_then_should_carry_canonical_ancestry_headers() {
979        let (record, conversation, source, correlation) = ids();
980        let mut envelope =
981            AgentEnvelope::command(record, conversation, source, correlation, b"x".to_vec());
982        envelope.parent = Some(laser_wire::agent::ConversationId::from_u128(2));
983        envelope.root = Some(laser_wire::agent::ConversationId::from_u128(1));
984        let headers = agdx_headers(&envelope, ContentType::Cbor, false).expect("headers stamp");
985        for (name, expected) in [
986            (PARENT_CONVERSATION_ID, envelope.parent.expect("parent")),
987            (ROOT_CONVERSATION_ID, envelope.root.expect("root")),
988        ] {
989            let value = headers
990                .get(&HeaderKey::from_str(name).expect("key"))
991                .expect("header");
992            assert_eq!(value.as_str().expect("string"), expected.to_string());
993        }
994    }
995
996    // The payload half of record conformance: the verb path's encoding is
997    // byte-identical to the shared golden corpus.
998    #[test]
999    fn given_the_canonical_command_when_encoded_then_should_match_the_corpus() {
1000        let (record, conversation, source, correlation) = ids();
1001        let command = AgentEnvelope::command(
1002            record,
1003            conversation,
1004            source,
1005            correlation,
1006            br#"{"ask":"plan the rollout"}"#.to_vec(),
1007        )
1008        .with_target("target-agent".parse().expect("valid agent id"))
1009        .with_idempotency_key("job-123-attempt-2".parse().expect("valid key"))
1010        .with_deadline_micros(1_717_171_777_000_000)
1011        .with_operation(OPERATION_CHAT)
1012        .with_metadata("priority", "high");
1013        let encoded = encode_named(&command).expect("encodes");
1014        assert_matches("agent_command.bin", &encoded);
1015    }
1016
1017    #[test]
1018    fn given_an_oversized_chunk_body_when_checked_then_should_reject_before_publish() {
1019        let body = vec![0u8; MAX_CHUNK_BODY_BYTES + 1];
1020        let error = ensure_chunk_body_within_cap(&body).expect_err("oversized body rejects");
1021        assert!(matches!(error, LaserError::Invalid(_)));
1022    }
1023
1024    #[test]
1025    fn given_the_canonical_record_when_encoded_then_should_match_record_fixture() {
1026        let (record, conversation, source, correlation) = ids();
1027        let envelope = AgentEnvelope::command(
1028            record,
1029            conversation,
1030            source,
1031            correlation,
1032            br#"{"ask":"plan the rollout"}"#.to_vec(),
1033        )
1034        .with_target("target-agent".parse().expect("valid agent id"))
1035        .with_idempotency_key("job-123-attempt-2".parse().expect("valid key"))
1036        .with_deadline_micros(1_717_171_777_000_000)
1037        .with_operation(OPERATION_CHAT)
1038        .with_metadata("priority", "high");
1039        let headers = agdx_headers(&envelope, ContentType::Json, false).expect("headers stamp");
1040        let payload = encode_named(&envelope).expect("payload encodes");
1041        let record =
1042            CanonicalAlpRecord::from_headers(envelope.conversation.to_string(), headers, payload);
1043        let encoded = encode_named(&record).expect("record fixture encodes");
1044        assert_record_fixture("agent_record.bin", &encoded);
1045    }
1046
1047    #[derive(Serialize)]
1048    struct CanonicalAlpRecord {
1049        partition_key: String,
1050        headers: BTreeMap<String, CanonicalHeader>,
1051        payload: Binary,
1052    }
1053
1054    impl CanonicalAlpRecord {
1055        fn from_headers(
1056            partition_key: String,
1057            headers: BTreeMap<HeaderKey, HeaderValue>,
1058            payload: Vec<u8>,
1059        ) -> Self {
1060            let mut canonical_headers = BTreeMap::new();
1061            for (key, value) in headers {
1062                let key = key.as_str().expect("header key is utf8").to_owned();
1063                let kind = match key.as_str() {
1064                    AGENT_VERSION => "u32",
1065                    CONTENT_TYPE => "u8",
1066                    CONVERSATION_ID | TARGET_AGENT_ID | AGENT_ID => "string",
1067                    _ => "bytes",
1068                };
1069                canonical_headers.insert(
1070                    key,
1071                    CanonicalHeader {
1072                        kind,
1073                        bytes: Binary(value.as_bytes().to_vec()),
1074                    },
1075                );
1076            }
1077            Self {
1078                partition_key,
1079                headers: canonical_headers,
1080                payload: Binary(payload),
1081            }
1082        }
1083    }
1084
1085    #[derive(Serialize)]
1086    struct CanonicalHeader {
1087        kind: &'static str,
1088        bytes: Binary,
1089    }
1090
1091    struct Binary(Vec<u8>);
1092
1093    impl Serialize for Binary {
1094        fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
1095        where
1096            S: Serializer,
1097        {
1098            serializer.serialize_bytes(&self.0)
1099        }
1100    }
1101
1102    fn assert_record_fixture(name: &str, encoded: &[u8]) {
1103        let path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
1104            .parent()
1105            .expect("sdk crate has workspace parent")
1106            .join("wire")
1107            .join("fixtures")
1108            .join(name);
1109        if env::var("AGDX_WIRE_FIXTURES_REGEN").is_ok() {
1110            fs::write(&path, encoded).expect("write record fixture");
1111        }
1112        let golden = fs::read(&path).unwrap_or_else(|error| {
1113            panic!("read fixture {name}: {error} (regen with AGDX_WIRE_FIXTURES_REGEN=1)")
1114        });
1115        assert_eq!(
1116            encoded, golden,
1117            "fixture `{name}` drifted from the canonical AGDX record"
1118        );
1119    }
1120
1121    fn ids() -> (RecordId, ConversationId, AgentId, CorrelationId) {
1122        (
1123            RecordId::from_u128(0x0190_3c1f_aa00_0000_0000_0000_0000_0001),
1124            ConversationId::from_u128(0x0190_3c1f_aa00_0000_0000_0000_0000_0002),
1125            "source-agent".parse().expect("valid agent id"),
1126            CorrelationId::from_u128(0x0190_3c1f_aa00_0000_0000_0000_0000_0005),
1127        )
1128    }
1129}
laser_sdk/agent/agdx.rs

laser_sdk/agent/
agdx.rs
