laser_sdk/provenance/
runtime.rs

1use super::keys;
2use crate::types::{AgentId, ConversationId, MessageId};
3use iggy::prelude::{HeaderKey, HeaderKind, HeaderValue, IggyError, IggyMessage, IggyTimestamp};
4use laser_wire::headers::{Addressee, HeaderField, RecordHeaders};
5use std::collections::BTreeMap;
6use std::str::FromStr;
7
8pub use super::topic::AgentTopic;
9pub use laser_wire::headers::{HEADER_FRAMING_BYTES, HEADER_SOFT_CAP, HEADER_VALUE_MAX};
10
11/// Why encoding or decoding provenance headers failed.
12#[derive(Debug, thiserror::Error)]
13pub enum ProvenanceError {
14    #[error("missing required header `{0}`")]
15    MissingRequired(&'static str),
16    #[error("provenance headers {got}B exceed soft cap {cap}B")]
17    TooLarge { got: usize, cap: usize },
18    #[error("invalid value for header `{0}`")]
19    InvalidValue(&'static str),
20    #[error("header `{key}` value must not contain control characters or NUL")]
21    InvalidValueBytes { key: &'static str },
22    #[error("non-finite floating-point value for header `{0}`")]
23    NonFinite(&'static str),
24    #[error("header `{0}` value must not be empty")]
25    EmptyValue(&'static str),
26    #[error("header `{key}` value is {got}B, exceeds max {max}B")]
27    ValueTooLong {
28        key: &'static str,
29        got: usize,
30        max: usize,
31    },
32    #[error("malformed Iggy headers on this message: {0}")]
33    MalformedHeaders(String),
34    #[error(transparent)]
35    Header(#[from] IggyError),
36    #[error(transparent)]
37    Id(#[from] crate::types::IdError),
38}
39
40/// Token / cost usage for an LLM call, carried on provenance for rollup.
41#[derive(Debug, Clone, Default, PartialEq, bon::Builder)]
42pub struct LlmUsage {
43    /// Prompt tokens, if known.
44    pub input_tokens: Option<u64>,
45    /// Completion tokens, if known.
46    pub output_tokens: Option<u64>,
47    /// Call cost in USD, if known.
48    pub cost_usd: Option<f64>,
49}
50
51/// The agentic message spine: conversation, causality, routing, usage. Encoded to/from message headers.
52#[derive(Debug, Clone, bon::Builder)]
53pub struct Provenance {
54    /// The conversation this message belongs to (the partition key).
55    pub conversation_id: ConversationId,
56    /// The message this one is a reply to, if any.
57    pub causal_parent: Option<MessageId>,
58    /// The conversation this was spawned from (sub-conversations).
59    pub parent_conversation_id: Option<ConversationId>,
60    /// The root of the conversation tree (sub-conversations).
61    pub root_conversation_id: Option<ConversationId>,
62    /// The agent that produced this message.
63    pub agent: Option<AgentId>,
64    /// The agent this message is addressed to (set by `Router`).
65    pub target_agent_id: Option<AgentId>,
66    /// LLM token / cost usage for this step.
67    pub usage: Option<LlmUsage>,
68    /// Drop-dead time. A consumer past it dead-letters the message.
69    pub deadline: Option<IggyTimestamp>,
70    /// Business dedup key: a producer-supplied key that suppresses a duplicate of
71    /// the same logical operation. Distinct from [`correlation_id`](Self::correlation_id):
72    /// setting a real dedup key and retrying no longer cross-matches replies.
73    pub idempotency_key: Option<String>,
74    /// Reply-correlation key: matches a reply to the request that minted it,
75    /// carried in the `agdx.corr` header. [`Laser::request`](crate::laser::Laser)
76    /// mints one when unset. Kept separate from [`idempotency_key`](Self::idempotency_key)
77    /// so a business dedup key and a reply correlation never overload one slot.
78    pub correlation_id: Option<String>,
79    /// The strictly-monotonic per-task fence the producer held, when this message
80    /// carries a fenced effect. The reliable consumer drops a message whose fence
81    /// is below the highest it has accepted for the task (a stale-holder replay).
82    pub fence_token: Option<u64>,
83}
84
85impl Provenance {
86    /// The Iggy partition key (the conversation id), so one conversation stays ordered.
87    pub fn partition_key(&self) -> String {
88        self.conversation_id.to_string()
89    }
90}
91
92impl TryFrom<&Provenance> for BTreeMap<HeaderKey, HeaderValue> {
93    type Error = ProvenanceError;
94
95    fn try_from(p: &Provenance) -> Result<Self, Self::Error> {
96        let mut headers = RecordHeaders::new(p.conversation_id.into());
97        headers.parent = p.parent_conversation_id.map(Into::into);
98        headers.root = p.root_conversation_id.map(Into::into);
99        headers.agent = p.agent.as_ref().map(wire_agent).transpose()?;
100        headers.addressee = p
101            .target_agent_id
102            .as_ref()
103            .map(|target| wire_agent(target).map(Addressee::Agent))
104            .transpose()?;
105        headers.causal_parent = p.causal_parent.map(|parent| parent.to_string());
106        headers.idempotency_key = p.idempotency_key.clone();
107        headers.correlation = p.correlation_id.clone();
108        headers.fence = p.fence_token;
109        headers.deadline_micros = p.deadline.map(|deadline| deadline.as_micros());
110        if let Some(usage) = &p.usage {
111            headers.input_tokens = usage.input_tokens;
112            headers.output_tokens = usage.output_tokens;
113            headers.cost_usd = usage.cost_usd;
114        }
115        header_map(&headers)
116    }
117}
118
119/// The Iggy header map of one record's header block, checked against the
120/// value and soft size caps.
121pub(crate) fn header_map(
122    headers: &RecordHeaders,
123) -> Result<BTreeMap<HeaderKey, HeaderValue>, ProvenanceError> {
124    let mut map = BTreeMap::new();
125    for (key, field) in headers.encode() {
126        let value = match field {
127            HeaderField::Text(text) => text_value(key, &text)?,
128            HeaderField::Uint8(value) => HeaderValue::from(value),
129            HeaderField::Uint32(value) => HeaderValue::from(value),
130            HeaderField::Uint64(value) => HeaderValue::from(value),
131            HeaderField::Float64(value) if value.is_finite() => HeaderValue::from(value),
132            HeaderField::Float64(_) => return Err(ProvenanceError::NonFinite(key)),
133        };
134        map.insert(HeaderKey::from_str(key)?, value);
135    }
136    let size: usize = map
137        .iter()
138        .map(|(k, v)| k.as_bytes().len() + v.as_bytes().len() + HEADER_FRAMING_BYTES)
139        .sum();
140    if size > HEADER_SOFT_CAP {
141        return Err(ProvenanceError::TooLarge {
142            got: size,
143            cap: HEADER_SOFT_CAP,
144        });
145    }
146    Ok(map)
147}
148
149fn wire_agent(agent: &AgentId) -> Result<laser_wire::agent::AgentId, ProvenanceError> {
150    agent
151        .as_str()
152        .parse()
153        .map_err(|_| ProvenanceError::InvalidValue(keys::AGENT_ID))
154}
155
156impl TryFrom<&IggyMessage> for Provenance {
157    type Error = ProvenanceError;
158
159    fn try_from(message: &IggyMessage) -> Result<Self, Self::Error> {
160        let headers = message
161            .user_headers_map()
162            .map_err(|err| ProvenanceError::MalformedHeaders(err.to_string()))?
163            .unwrap_or_default();
164        provenance_from_headers(&headers)
165    }
166}
167
168/// Decode provenance from an already-parsed Iggy header map. Split out so a
169/// caller that has the map in hand (the envelope-aware consumer, which checks
170/// for `agdx.av` first) decodes the headers once rather than re-parsing them.
171pub(crate) fn provenance_from_headers(
172    headers: &BTreeMap<HeaderKey, HeaderValue>,
173) -> Result<Provenance, ProvenanceError> {
174    let mut conversation_id: Option<ConversationId> = None;
175    let mut causal_parent: Option<MessageId> = None;
176    let mut parent_conversation_id: Option<ConversationId> = None;
177    let mut root_conversation_id: Option<ConversationId> = None;
178    let mut agent: Option<AgentId> = None;
179    let mut target_agent_id: Option<AgentId> = None;
180    let mut idempotency_key: Option<String> = None;
181    let mut correlation_id: Option<String> = None;
182    let mut fence_token: Option<u64> = None;
183    let mut deadline: Option<IggyTimestamp> = None;
184    let mut usage = LlmUsage::default();
185    let mut has_usage = false;
186
187    for (key, value) in headers {
188        let key_str = key
189            .as_str()
190            .map_err(|err| ProvenanceError::MalformedHeaders(err.to_string()))?;
191        // Match on the key first. A header outside the provenance dictionary
192        // (the AGDX `agdx.ct` u8, `agdx.av` u32, any app-custom key) is foreign
193        // and ignored. Ids ride as strings and numbers ride typed. A known key
194        // with another value kind is corruption, not something to drop, so it
195        // is a decode error.
196        match key_str {
197            keys::CONVERSATION_ID => {
198                conversation_id = Some(match value.kind() {
199                    HeaderKind::Uint128 => laser_wire::agent::ConversationId::from_u128(
200                        value
201                            .as_uint128()
202                            .map_err(|_| ProvenanceError::InvalidValue(keys::CONVERSATION_ID))?,
203                    )
204                    .into(),
205                    _ => str_value(value, keys::CONVERSATION_ID)?.parse()?,
206                });
207            }
208            keys::CAUSAL_PARENT => {
209                causal_parent = Some(str_value(value, keys::CAUSAL_PARENT)?.parse()?);
210            }
211            keys::PARENT_CONVERSATION_ID => {
212                parent_conversation_id =
213                    Some(str_value(value, keys::PARENT_CONVERSATION_ID)?.parse()?);
214            }
215            keys::ROOT_CONVERSATION_ID => {
216                root_conversation_id = Some(str_value(value, keys::ROOT_CONVERSATION_ID)?.parse()?);
217            }
218            keys::AGENT_ID => agent = Some(str_value(value, keys::AGENT_ID)?.parse()?),
219            keys::TARGET_AGENT_ID => {
220                // Broadcast `*` addresses every agent and is never an agent id.
221                let target = str_value(value, keys::TARGET_AGENT_ID)?;
222                if target != laser_wire::headers::BROADCAST {
223                    target_agent_id = Some(target.parse()?);
224                }
225            }
226            keys::IDEMPOTENCY_KEY => {
227                idempotency_key = Some(str_value(value, keys::IDEMPOTENCY_KEY)?.to_owned());
228            }
229            keys::CORRELATION_ID => {
230                correlation_id = Some(str_value(value, keys::CORRELATION_ID)?.to_owned());
231            }
232            // A malformed fence must be a decode error, never `.ok()`-ed to `None`:
233            // silently dropping it would skip the gate and let a tampered record
234            // through as unfenced.
235            keys::FENCE => fence_token = Some(u64_value(value, keys::FENCE)?),
236            keys::DEADLINE => {
237                deadline = Some(IggyTimestamp::from(u64_value(value, keys::DEADLINE)?));
238            }
239            keys::USAGE_INPUT_TOKENS => {
240                usage.input_tokens = Some(u64_value(value, keys::USAGE_INPUT_TOKENS)?);
241                has_usage = true;
242            }
243            keys::USAGE_OUTPUT_TOKENS => {
244                usage.output_tokens = Some(u64_value(value, keys::USAGE_OUTPUT_TOKENS)?);
245                has_usage = true;
246            }
247            keys::COST_USD => {
248                let cost = value
249                    .as_float64()
250                    .map_err(|_| ProvenanceError::InvalidValue(keys::COST_USD))?;
251                if !cost.is_finite() {
252                    return Err(ProvenanceError::NonFinite(keys::COST_USD));
253                }
254                usage.cost_usd = Some(cost);
255                has_usage = true;
256            }
257            _ => {}
258        }
259    }
260
261    Ok(Provenance {
262        conversation_id: conversation_id
263            .ok_or(ProvenanceError::MissingRequired(keys::CONVERSATION_ID))?,
264        causal_parent,
265        parent_conversation_id,
266        root_conversation_id,
267        agent,
268        target_agent_id,
269        usage: has_usage.then_some(usage),
270        deadline,
271        idempotency_key,
272        correlation_id,
273        fence_token,
274    })
275}
276
277fn text_value(key: &'static str, value: &str) -> Result<HeaderValue, ProvenanceError> {
278    if value.is_empty() {
279        return Err(ProvenanceError::EmptyValue(key));
280    }
281    if value.len() > HEADER_VALUE_MAX {
282        return Err(ProvenanceError::ValueTooLong {
283            key,
284            got: value.len(),
285            max: HEADER_VALUE_MAX,
286        });
287    }
288    // Reject ASCII control characters and DEL. Header values ride through Iggy
289    // and downstream log formatters / OTel exporters, and a stray `\n` or `\0`
290    // breaks line-oriented parsers and corrupts traces.
291    if value.bytes().any(|b| b < 0x20 || b == 0x7f) {
292        return Err(ProvenanceError::InvalidValueBytes { key });
293    }
294    Ok(HeaderValue::from_str(value)?)
295}
296
297// A malformed number must be a decode error, never dropped: a silently missing
298// fence would skip the gate and let a tampered record through as unfenced.
299fn u64_value(value: &HeaderValue, key: &'static str) -> Result<u64, ProvenanceError> {
300    value
301        .as_uint64()
302        .map_err(|_| ProvenanceError::InvalidValue(key))
303}
304
305// A known provenance key's value must be a string, and a non-string value there is
306// corruption, reported as an error rather than silently dropped.
307fn str_value<'a>(value: &'a HeaderValue, key: &'static str) -> Result<&'a str, ProvenanceError> {
308    value
309        .as_str()
310        .map_err(|_| ProvenanceError::InvalidValue(key))
311}
312
313#[cfg(test)]
314mod tests {
315    use super::*;
316    use bytes::Bytes;
317
318    #[test]
319    fn given_provenance_when_round_tripped_through_headers_then_should_preserve_every_field() {
320        let conversation_id = ConversationId::new();
321        let provenance = Provenance::builder()
322            .conversation_id(conversation_id)
323            .causal_parent(MessageId::new(2, 7))
324            .agent("planner".parse().expect("planner is a valid agent id"))
325            .target_agent_id("executor".parse().expect("executor is a valid agent id"))
326            .idempotency_key("key-1".to_owned())
327            .correlation_id("corr-1".to_owned())
328            .fence_token(7)
329            .usage(
330                LlmUsage::builder()
331                    .input_tokens(10)
332                    .output_tokens(20)
333                    .build(),
334            )
335            .build();
336
337        let headers: BTreeMap<HeaderKey, HeaderValue> = (&provenance)
338            .try_into()
339            .expect("provenance should encode into headers");
340        let message = IggyMessage::builder()
341            .payload(Bytes::from_static(b"hello"))
342            .user_headers(headers)
343            .build()
344            .expect("the message should build");
345
346        let back =
347            Provenance::try_from(&message).expect("provenance should decode from the message");
348        assert_eq!(back.conversation_id, conversation_id);
349        assert_eq!(back.causal_parent, Some(MessageId::new(2, 7)));
350        assert_eq!(back.agent.expect("agent should be set").as_str(), "planner");
351        assert_eq!(
352            back.target_agent_id
353                .expect("target agent should be set")
354                .as_str(),
355            "executor"
356        );
357        assert_eq!(back.idempotency_key.as_deref(), Some("key-1"));
358        assert_eq!(back.correlation_id.as_deref(), Some("corr-1"));
359        assert_eq!(back.fence_token, Some(7));
360        let usage = back.usage.expect("usage should be set");
361        assert_eq!(usage.input_tokens, Some(10));
362        assert_eq!(usage.output_tokens, Some(20));
363    }
364
365    #[test]
366    fn given_a_malformed_fence_header_when_decoded_then_should_error_not_skip() {
367        let conversation = ConversationId::new();
368        let mut headers = BTreeMap::new();
369        headers.insert(
370            HeaderKey::from_str(keys::CONVERSATION_ID).expect("a valid header key"),
371            HeaderValue::from_str(&conversation.to_string()).expect("a valid header value"),
372        );
373        // A non-numeric fence must be rejected, never `.ok()`-ed to absent, or a
374        // tampered record would slip through the gate as unfenced.
375        headers.insert(
376            HeaderKey::from_str(keys::FENCE).expect("a valid header key"),
377            HeaderValue::from_str("not-a-number").expect("a valid header value"),
378        );
379        assert!(provenance_from_headers(&headers).is_err());
380    }
381
382    #[test]
383    fn given_a_message_without_a_conversation_id_when_decoded_then_should_error() {
384        let message = IggyMessage::builder()
385            .payload(Bytes::from_static(b"x"))
386            .build()
387            .expect("the message should build");
388        assert!(matches!(
389            Provenance::try_from(&message),
390            Err(ProvenanceError::MissingRequired(keys::CONVERSATION_ID))
391        ));
392    }
393
394    #[test]
395    fn given_a_typed_conversation_header_when_decoded_then_should_preserve_the_id() {
396        let conversation = ConversationId::new();
397        let mut headers = BTreeMap::new();
398        headers.insert(
399            HeaderKey::from_str(keys::CONVERSATION_ID).expect("a valid header key"),
400            HeaderValue::from(conversation.as_u128()),
401        );
402        let message = IggyMessage::builder()
403            .payload(Bytes::from_static(b"typed"))
404            .user_headers(headers)
405            .build()
406            .expect("the message should build");
407
408        let decoded = Provenance::try_from(&message).expect("typed conversation should decode");
409        assert_eq!(decoded.conversation_id, conversation);
410    }
411
412    #[test]
413    fn given_a_typed_non_string_header_when_decoded_then_should_skip_it_not_error() {
414        let conversation = ConversationId::new();
415        let mut headers = BTreeMap::new();
416        headers.insert(
417            HeaderKey::from_str(keys::CONVERSATION_ID).expect("a valid header key"),
418            HeaderValue::from_str(&conversation.to_string()).expect("a valid header value"),
419        );
420        // A typed (u8) header riding alongside provenance, like the AGDX `agdx.ct`.
421        headers.insert(
422            HeaderKey::from_str("agdx.ct").expect("a valid header key"),
423            HeaderValue::from(7u8),
424        );
425        let message = IggyMessage::builder()
426            .payload(Bytes::from_static(b"x"))
427            .user_headers(headers)
428            .build()
429            .expect("the message should build");
430        let provenance = Provenance::try_from(&message)
431            .expect("a foreign typed header is ignored, not a decode error");
432        assert_eq!(provenance.conversation_id, conversation);
433    }
434
435    #[test]
436    fn given_a_known_key_with_the_wrong_value_kind_when_decoded_then_should_error() {
437        let conversation = ConversationId::new();
438        let mut headers = BTreeMap::new();
439        headers.insert(
440            HeaderKey::from_str(keys::CONVERSATION_ID).expect("a valid header key"),
441            HeaderValue::from_str(&conversation.to_string()).expect("a valid header value"),
442        );
443        // `agdx.deadline` is a known provenance key that rides as a typed u64, so
444        // a text value there is corruption, not a header to silently drop.
445        headers.insert(
446            HeaderKey::from_str(keys::DEADLINE).expect("a valid header key"),
447            HeaderValue::from_str("42").expect("a valid header value"),
448        );
449        let message = IggyMessage::builder()
450            .payload(Bytes::from_static(b"x"))
451            .user_headers(headers)
452            .build()
453            .expect("the message should build");
454        assert!(matches!(
455            Provenance::try_from(&message),
456            Err(ProvenanceError::InvalidValue(keys::DEADLINE))
457        ));
458    }
459
460    #[test]
461    fn given_an_oversized_idempotency_key_when_encoded_then_should_report_a_clear_error() {
462        let provenance = Provenance::builder()
463            .conversation_id(ConversationId::new())
464            .idempotency_key("x".repeat(HEADER_VALUE_MAX + 1))
465            .build();
466        let result: Result<BTreeMap<HeaderKey, HeaderValue>, _> = (&provenance).try_into();
467        assert!(matches!(
468            result,
469            Err(ProvenanceError::ValueTooLong {
470                key: keys::IDEMPOTENCY_KEY,
471                ..
472            })
473        ));
474    }
475
476    #[test]
477    fn given_an_empty_idempotency_key_when_encoded_then_should_report_a_clear_error() {
478        let provenance = Provenance::builder()
479            .conversation_id(ConversationId::new())
480            .idempotency_key(String::new())
481            .build();
482        let result: Result<BTreeMap<HeaderKey, HeaderValue>, _> = (&provenance).try_into();
483        assert!(matches!(
484            result,
485            Err(ProvenanceError::EmptyValue(keys::IDEMPOTENCY_KEY))
486        ));
487    }
488
489    #[test]
490    fn given_the_current_header_names_when_checked_then_they_stay_current() {
491        assert_eq!(keys::CONVERSATION_ID, "gen_ai.conversation.id");
492        assert_eq!(keys::AGENT_ID, "gen_ai.agent.id");
493        assert_eq!(keys::USAGE_INPUT_TOKENS, "gen_ai.usage.input_tokens");
494        assert_eq!(keys::USAGE_OUTPUT_TOKENS, "gen_ai.usage.output_tokens");
495        assert_eq!(keys::CAUSAL_PARENT, "agdx.cause");
496        assert_eq!(keys::PARENT_CONVERSATION_ID, "agdx.parent_conv");
497        assert_eq!(keys::ROOT_CONVERSATION_ID, "agdx.root_conv");
498        assert_eq!(keys::TARGET_AGENT_ID, "agdx.to");
499        assert_eq!(keys::IDEMPOTENCY_KEY, "agdx.idem");
500        assert_eq!(keys::COST_USD, "agdx.cost");
501    }
502}
laser_sdk/provenance/runtime.rs

laser_sdk/provenance/
runtime.rs
