lean_ctx/core/
capsule_transport.rs

1//! Local reference transport for authenticated OCLA A2A capsule delivery (P11).
2//!
3//! In-process delivery queue that accepts signed, payload-free capsule manifests.
4//! Bytes are measured exactly after serialization; the token count is a local
5//! tokenizer proxy. This is delivery evidence only — never convertible into
6//! compression savings or provider billing.
7
8use std::collections::{BTreeMap, VecDeque};
9use std::sync::Mutex;
10use std::time::{Duration, Instant};
11
12use serde::{Deserialize, Serialize};
13
14use crate::core::context_capsule::SignedContextCapsuleV1;
15
16const MAX_INBOX_SIZE: usize = 128;
17const MAX_DELIVERY_ATTEMPTS: u8 = 3;
18const RESERVATION_TTL: Duration = Duration::from_mins(1);
19
20#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
21#[serde(rename_all = "snake_case")]
22pub enum DeliveredTokenCountKindV1 {
23    LocalTokenizerProxy,
24}
25
26/// Delivery receipt measured after a concrete local queue write.
27#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
28pub struct AgentRelayDeliveryV1 {
29    pub relay_id: String,
30    pub capsule_ref: String,
31    pub recipient_agent_id: String,
32    pub delivered_bytes: u64,
33    pub delivered_tokens: u64,
34    pub token_count_kind: DeliveredTokenCountKindV1,
35}
36
37#[derive(Clone)]
38struct QueuedCapsule {
39    signed_json: String,
40    delivery: AgentRelayDeliveryV1,
41    reservation_id: Option<String>,
42    reserved_at: Option<Instant>,
43    delivery_attempts: u8,
44}
45
46/// Bounded local transport. Deployments must place a durable, authenticated
47/// transport behind the same receipt semantics before remote use.
48pub struct LocalSignedCapsuleTransport {
49    inboxes: Mutex<BTreeMap<String, VecDeque<QueuedCapsule>>>,
50    dead_letters: Mutex<BTreeMap<String, VecDeque<QueuedCapsule>>>,
51}
52
53impl Default for LocalSignedCapsuleTransport {
54    fn default() -> Self {
55        Self {
56            inboxes: Mutex::new(BTreeMap::new()),
57            dead_letters: Mutex::new(BTreeMap::new()),
58        }
59    }
60}
61
62impl LocalSignedCapsuleTransport {
63    /// Deliver a signed capsule to a recipient's inbox.
64    ///
65    /// The signature is verified against the caller-pinned `signer` key before
66    /// anything is queued; a capsule whose signature does not verify, or that
67    /// was signed by another key, is rejected.
68    pub fn deliver(
69        &self,
70        signed: &SignedContextCapsuleV1,
71        recipient_agent_id: &str,
72        signer: &ed25519_dalek::VerifyingKey,
73    ) -> Result<AgentRelayDeliveryV1, TransportError> {
74        signed
75            .verify(signer)
76            .map_err(|e| TransportError::Validation(e.to_string()))?;
77
78        let envelope = signed
79            .capsule
80            .agent_envelope(recipient_agent_id)
81            .map_err(|e| TransportError::Validation(e.to_string()))?;
82
83        let signed_json = serde_json::to_string(signed)
84            .map_err(|e| TransportError::Serialization(e.to_string()))?;
85
86        let delivered_bytes = u64::try_from(signed_json.len()).unwrap_or(u64::MAX);
87        let delivered_tokens = estimate_tokens(&signed_json);
88
89        let delivery = AgentRelayDeliveryV1 {
90            relay_id: envelope.relay_id.clone(),
91            capsule_ref: envelope.capsule_ref,
92            recipient_agent_id: recipient_agent_id.to_string(),
93            delivered_bytes,
94            delivered_tokens,
95            token_count_kind: DeliveredTokenCountKindV1::LocalTokenizerProxy,
96        };
97
98        let mut inboxes = self
99            .inboxes
100            .lock()
101            .unwrap_or_else(std::sync::PoisonError::into_inner);
102        let inbox = inboxes.entry(recipient_agent_id.to_string()).or_default();
103
104        let dead_letter_depth = self
105            .dead_letters
106            .lock()
107            .unwrap_or_else(std::sync::PoisonError::into_inner)
108            .get(recipient_agent_id)
109            .map_or(0, VecDeque::len);
110        if inbox.len().saturating_add(dead_letter_depth) >= MAX_INBOX_SIZE {
111            return Err(TransportError::Capacity);
112        }
113
114        inbox.push_back(QueuedCapsule {
115            signed_json,
116            delivery: delivery.clone(),
117            reservation_id: None,
118            reserved_at: None,
119            delivery_attempts: 0,
120        });
121
122        Ok(delivery)
123    }
124
125    /// Receive the next capsule from an agent's inbox.
126    pub fn receive(
127        &self,
128        agent_id: &str,
129    ) -> Option<(SignedContextCapsuleV1, AgentRelayDeliveryV1)> {
130        let mut inboxes = self
131            .inboxes
132            .lock()
133            .unwrap_or_else(std::sync::PoisonError::into_inner);
134        let inbox = inboxes.get_mut(agent_id)?;
135        if inbox.front()?.reservation_id.is_some() {
136            return None;
137        }
138        let queued = inbox.pop_front()?;
139        let signed: SignedContextCapsuleV1 = serde_json::from_str(&queued.signed_json).ok()?;
140        Some((signed, queued.delivery))
141    }
142
143    /// Exclusively reserve the next capsule until acknowledged or negatively acknowledged.
144    pub fn reserve(
145        &self,
146        agent_id: &str,
147        reservation_id: &str,
148    ) -> Result<Option<(SignedContextCapsuleV1, AgentRelayDeliveryV1)>, TransportError> {
149        self.reserve_at(agent_id, reservation_id, Instant::now())
150    }
151
152    fn reserve_at(
153        &self,
154        agent_id: &str,
155        reservation_id: &str,
156        now: Instant,
157    ) -> Result<Option<(SignedContextCapsuleV1, AgentRelayDeliveryV1)>, TransportError> {
158        if reservation_id.is_empty() || reservation_id.len() > 256 {
159            return Err(TransportError::Validation("invalid reservation id".into()));
160        }
161        let mut inboxes = self
162            .inboxes
163            .lock()
164            .unwrap_or_else(std::sync::PoisonError::into_inner);
165        let Some(inbox) = inboxes.get_mut(agent_id) else {
166            return Ok(None);
167        };
168        let expired_exhausted = inbox.front().is_some_and(|front| {
169            front.delivery_attempts >= MAX_DELIVERY_ATTEMPTS
170                && front.reserved_at.is_some_and(|reserved_at| {
171                    now.saturating_duration_since(reserved_at) >= RESERVATION_TTL
172                })
173        });
174        if expired_exhausted {
175            let exhausted = inbox.pop_front().expect("front checked above");
176            self.dead_letters
177                .lock()
178                .unwrap_or_else(std::sync::PoisonError::into_inner)
179                .entry(agent_id.to_string())
180                .or_default()
181                .push_back(exhausted);
182            return Ok(None);
183        }
184        let Some(front) = inbox.front_mut() else {
185            return Ok(None);
186        };
187        if front.reserved_at.is_some_and(|reserved_at| {
188            now.saturating_duration_since(reserved_at) >= RESERVATION_TTL
189        }) {
190            front.reservation_id = None;
191            front.reserved_at = None;
192        }
193        if front.reservation_id.is_some() {
194            return Ok(None);
195        }
196        let signed: SignedContextCapsuleV1 = match serde_json::from_str(&front.signed_json) {
197            Ok(signed) => signed,
198            Err(error) => {
199                let poisoned = inbox.pop_front().expect("front checked above");
200                self.dead_letters
201                    .lock()
202                    .unwrap_or_else(std::sync::PoisonError::into_inner)
203                    .entry(agent_id.to_string())
204                    .or_default()
205                    .push_back(poisoned);
206                return Err(TransportError::Serialization(error.to_string()));
207            }
208        };
209        front.reservation_id = Some(reservation_id.to_string());
210        front.reserved_at = Some(now);
211        front.delivery_attempts = front.delivery_attempts.saturating_add(1);
212        Ok(Some((signed, front.delivery.clone())))
213    }
214
215    pub fn acknowledge(&self, agent_id: &str, relay_id: &str, reservation_id: &str) -> bool {
216        let mut inboxes = self
217            .inboxes
218            .lock()
219            .unwrap_or_else(std::sync::PoisonError::into_inner);
220        let Some(inbox) = inboxes.get_mut(agent_id) else {
221            return false;
222        };
223        if inbox.front().is_some_and(|queued| {
224            queued.delivery.relay_id == relay_id
225                && queued.reservation_id.as_deref() == Some(reservation_id)
226                && queued.reserved_at.is_some_and(|reserved_at| {
227                    Instant::now().saturating_duration_since(reserved_at) < RESERVATION_TTL
228                })
229        }) {
230            inbox.pop_front();
231            return true;
232        }
233        false
234    }
235
236    /// Release a failed reservation; terminal retries move to a bounded dead-letter queue.
237    pub fn negative_acknowledge(
238        &self,
239        agent_id: &str,
240        relay_id: &str,
241        reservation_id: &str,
242    ) -> bool {
243        let mut inboxes = self
244            .inboxes
245            .lock()
246            .unwrap_or_else(std::sync::PoisonError::into_inner);
247        let Some(inbox) = inboxes.get_mut(agent_id) else {
248            return false;
249        };
250        let matches = inbox.front().is_some_and(|queued| {
251            queued.delivery.relay_id == relay_id
252                && queued.reservation_id.as_deref() == Some(reservation_id)
253                && queued.reserved_at.is_some_and(|reserved_at| {
254                    Instant::now().saturating_duration_since(reserved_at) < RESERVATION_TTL
255                })
256        });
257        if !matches {
258            return false;
259        }
260        if inbox
261            .front()
262            .is_some_and(|queued| queued.delivery_attempts >= MAX_DELIVERY_ATTEMPTS)
263        {
264            let exhausted = inbox.pop_front().expect("front checked above");
265            self.dead_letters
266                .lock()
267                .unwrap_or_else(std::sync::PoisonError::into_inner)
268                .entry(agent_id.to_string())
269                .or_default()
270                .push_back(exhausted);
271        } else if let Some(front) = inbox.front_mut() {
272            front.reservation_id = None;
273            front.reserved_at = None;
274        }
275        true
276    }
277
278    pub fn dead_letter_depth(&self, agent_id: &str) -> usize {
279        self.dead_letters
280            .lock()
281            .unwrap_or_else(std::sync::PoisonError::into_inner)
282            .get(agent_id)
283            .map_or(0, VecDeque::len)
284    }
285
286    /// Quarantine a delivery whose committed consumer lost its ACK invariant.
287    pub fn quarantine_reserved(
288        &self,
289        agent_id: &str,
290        relay_id: &str,
291        reservation_id: &str,
292    ) -> bool {
293        let mut inboxes = self
294            .inboxes
295            .lock()
296            .unwrap_or_else(std::sync::PoisonError::into_inner);
297        let Some(inbox) = inboxes.get_mut(agent_id) else {
298            return false;
299        };
300        let Some(index) = inbox.iter().position(|queued| {
301            queued.delivery.relay_id == relay_id
302                && queued.reservation_id.as_deref() == Some(reservation_id)
303        }) else {
304            return false;
305        };
306        let Some(queued) = inbox.remove(index) else {
307            return false;
308        };
309        self.dead_letters
310            .lock()
311            .unwrap_or_else(std::sync::PoisonError::into_inner)
312            .entry(agent_id.to_string())
313            .or_default()
314            .push_back(queued);
315        true
316    }
317
318    /// Peek at inbox depth for a given agent.
319    pub fn inbox_depth(&self, agent_id: &str) -> usize {
320        let inboxes = self
321            .inboxes
322            .lock()
323            .unwrap_or_else(std::sync::PoisonError::into_inner);
324        inboxes.get(agent_id).map_or(0, VecDeque::len)
325    }
326}
327
328/// Simple token estimation: ~4 chars per token (conservative for JSON).
329fn estimate_tokens(text: &str) -> u64 {
330    u64::try_from(text.len() / 4).unwrap_or(u64::MAX).max(1)
331}
332
333#[derive(Debug, thiserror::Error)]
334pub enum TransportError {
335    #[error("validation failed: {0}")]
336    Validation(String),
337    #[error("serialization failed: {0}")]
338    Serialization(String),
339    #[error("recipient capsule queue is at capacity")]
340    Capacity,
341}
342
343// ─── Tests ───────────────────────────────────────────────────────────────────
344
345#[cfg(test)]
346mod tests {
347    use super::*;
348    use crate::core::context_capsule::*;
349
350    fn test_signed_capsule() -> SignedContextCapsuleV1 {
351        let mut capsule = ContextCapsuleV1 {
352            schema_version: CONTEXT_CAPSULE_SCHEMA_VERSION,
353            capsule_id: "capsule:pending".to_string(),
354            request_id: "request:1".to_string(),
355            session_id: "session:1".to_string(),
356            agent_id: "sender-agent".to_string(),
357            intent_ref: "intent:test".to_string(),
358            task_ref: "task:1".to_string(),
359            expected_outcome_ref: "outcome:pass".to_string(),
360            acceptance_criteria_refs: vec!["criteria:ci-green".to_string()],
361            references: vec![ContextCapsuleReferenceV1 {
362                kind: CapsuleReferenceKindV1::File,
363                content_ref: "blake3:file1".to_string(),
364                freshness_ref: "freshness:1".to_string(),
365                recovery_ref: None,
366            }],
367            finding_refs: vec![],
368            decision_refs: vec![],
369            uncertainty_refs: vec![],
370            negative_result_refs: vec![],
371            source_ref: "source:test".to_string(),
372            policy_ref: "policy:default".to_string(),
373            contract_ref: "contract:v1".to_string(),
374            freshness_ref: "freshness:now".to_string(),
375            sensitivity: CapsuleSensitivityV1::Internal,
376            allowed_agent_ids: vec!["recipient-agent".to_string()],
377            budget: ContextCapsuleBudgetV1 {
378                tokens_used: 50,
379                tokens_remaining: 950,
380                cost_micros_used: 5,
381                cost_micros_remaining: 95,
382                latency_ms_used: 10,
383                latency_ms_remaining: 90,
384            },
385            chain: ContextCapsuleChainV1 {
386                chain_id: "chain:test".to_string(),
387                parent_capsule_ref: None,
388                owner_agent_id: "sender-agent".to_string(),
389                attribution_ref: "attribution:1".to_string(),
390                hop: 0,
391            },
392            quality_signal_refs: vec![],
393            recovery_refs: vec![],
394            delta_from: None,
395        };
396        capsule.assign_capsule_id().unwrap();
397        let keypair = ed25519_dalek::SigningKey::from_bytes(&[42u8; 32]);
398        SignedContextCapsuleV1::sign(&capsule, &keypair).unwrap()
399    }
400
401    fn test_signer() -> ed25519_dalek::VerifyingKey {
402        ed25519_dalek::SigningKey::from_bytes(&[42u8; 32]).verifying_key()
403    }
404
405    #[test]
406    fn rejects_forged_or_foreign_signature_before_queueing() {
407        let transport = LocalSignedCapsuleTransport::default();
408
409        let mut forged = test_signed_capsule();
410        forged.signature = "00".repeat(64);
411        assert!(matches!(
412            transport.deliver(&forged, "recipient-agent", &test_signer()),
413            Err(TransportError::Validation(_))
414        ));
415
416        let mut tampered = test_signed_capsule();
417        tampered.capsule.budget.tokens_remaining = 9_999;
418        tampered.capsule.assign_capsule_id().unwrap();
419        assert!(
420            transport
421                .deliver(&tampered, "recipient-agent", &test_signer())
422                .is_err()
423        );
424
425        let foreign = ed25519_dalek::SigningKey::from_bytes(&[7u8; 32]).verifying_key();
426        assert!(
427            transport
428                .deliver(&test_signed_capsule(), "recipient-agent", &foreign)
429                .is_err()
430        );
431        assert_eq!(transport.inbox_depth("recipient-agent"), 0);
432    }
433
434    #[test]
435    fn deliver_and_receive_roundtrip() {
436        let transport = LocalSignedCapsuleTransport::default();
437        let signed = test_signed_capsule();
438
439        let delivery = transport
440            .deliver(&signed, "recipient-agent", &test_signer())
441            .unwrap();
442        assert!(delivery.delivered_bytes > 0);
443        assert!(delivery.delivered_tokens > 0);
444        assert_eq!(delivery.recipient_agent_id, "recipient-agent");
445        assert_eq!(transport.inbox_depth("recipient-agent"), 1);
446
447        let (received, receipt) = transport.receive("recipient-agent").unwrap();
448        assert_eq!(received.capsule.capsule_id, signed.capsule.capsule_id);
449        assert_eq!(receipt.relay_id, delivery.relay_id);
450        assert_eq!(transport.inbox_depth("recipient-agent"), 0);
451    }
452
453    #[test]
454    fn receive_from_empty_inbox_returns_none() {
455        let transport = LocalSignedCapsuleTransport::default();
456        assert!(transport.receive("nobody").is_none());
457    }
458
459    #[test]
460    fn rejects_unauthorized_recipient() {
461        let transport = LocalSignedCapsuleTransport::default();
462        let signed = test_signed_capsule();
463        assert!(
464            transport
465                .deliver(&signed, "unauthorized-agent", &test_signer())
466                .is_err()
467        );
468    }
469
470    #[test]
471    fn reservation_is_exclusive_and_ack_is_fenced() {
472        let transport = LocalSignedCapsuleTransport::default();
473        let delivery = transport
474            .deliver(&test_signed_capsule(), "recipient-agent", &test_signer())
475            .unwrap();
476        assert!(
477            transport
478                .reserve("recipient-agent", "worker:1")
479                .unwrap()
480                .is_some()
481        );
482        assert!(
483            transport
484                .reserve("recipient-agent", "worker:1")
485                .unwrap()
486                .is_none()
487        );
488        assert!(
489            transport
490                .reserve("recipient-agent", "worker:2")
491                .unwrap()
492                .is_none()
493        );
494        assert!(!transport.acknowledge("recipient-agent", &delivery.relay_id, "worker:2"));
495        assert!(transport.acknowledge("recipient-agent", &delivery.relay_id, "worker:1"));
496        assert_eq!(transport.inbox_depth("recipient-agent"), 0);
497    }
498
499    #[test]
500    fn expired_reservation_is_recoverable_with_a_new_fence() {
501        let transport = LocalSignedCapsuleTransport::default();
502        let delivery = transport
503            .deliver(&test_signed_capsule(), "recipient-agent", &test_signer())
504            .unwrap();
505        let now = Instant::now();
506        assert!(
507            transport
508                .reserve_at("recipient-agent", "worker:stale", now)
509                .unwrap()
510                .is_some()
511        );
512        assert!(
513            transport
514                .reserve_at("recipient-agent", "worker:fresh", now + RESERVATION_TTL)
515                .unwrap()
516                .is_some()
517        );
518        assert!(!transport.quarantine_reserved(
519            "recipient-agent",
520            &delivery.relay_id,
521            "worker:stale"
522        ));
523        assert!(transport.acknowledge("recipient-agent", &delivery.relay_id, "worker:fresh"));
524    }
525
526    #[test]
527    fn legacy_receive_cannot_steal_a_reserved_capsule() {
528        let transport = LocalSignedCapsuleTransport::default();
529        let delivery = transport
530            .deliver(&test_signed_capsule(), "recipient-agent", &test_signer())
531            .unwrap();
532        transport
533            .reserve("recipient-agent", "worker:1")
534            .unwrap()
535            .unwrap();
536        assert!(transport.receive("recipient-agent").is_none());
537        assert!(transport.quarantine_reserved("recipient-agent", &delivery.relay_id, "worker:1"));
538        assert_eq!(transport.inbox_depth("recipient-agent"), 0);
539        assert_eq!(transport.dead_letter_depth("recipient-agent"), 1);
540    }
541
542    #[test]
543    fn exhausted_negative_ack_moves_capsule_to_dead_letter_queue() {
544        let transport = LocalSignedCapsuleTransport::default();
545        let delivery = transport
546            .deliver(&test_signed_capsule(), "recipient-agent", &test_signer())
547            .unwrap();
548        for attempt in 1..=MAX_DELIVERY_ATTEMPTS {
549            let reservation = format!("worker:{attempt}");
550            assert!(
551                transport
552                    .reserve("recipient-agent", &reservation)
553                    .unwrap()
554                    .is_some()
555            );
556            assert!(transport.negative_acknowledge(
557                "recipient-agent",
558                &delivery.relay_id,
559                &reservation
560            ));
561        }
562        assert_eq!(transport.inbox_depth("recipient-agent"), 0);
563        assert_eq!(transport.dead_letter_depth("recipient-agent"), 1);
564    }
565
566    #[test]
567    fn malformed_front_is_quarantined_instead_of_poisoning_inbox() {
568        let transport = LocalSignedCapsuleTransport::default();
569        transport
570            .deliver(&test_signed_capsule(), "recipient-agent", &test_signer())
571            .unwrap();
572        transport
573            .inboxes
574            .lock()
575            .unwrap()
576            .get_mut("recipient-agent")
577            .unwrap()
578            .front_mut()
579            .unwrap()
580            .signed_json = "{".into();
581
582        assert!(transport.reserve("recipient-agent", "worker:1").is_err());
583        assert_eq!(transport.inbox_depth("recipient-agent"), 0);
584        assert_eq!(transport.dead_letter_depth("recipient-agent"), 1);
585    }
586
587    #[test]
588    fn full_inbox_rejects_delivery_without_evicting_existing_capsules() {
589        let transport = LocalSignedCapsuleTransport::default();
590        let signed = test_signed_capsule();
591        for _ in 0..MAX_INBOX_SIZE {
592            transport
593                .deliver(&signed, "recipient-agent", &test_signer())
594                .unwrap();
595        }
596        assert!(matches!(
597            transport.deliver(&signed, "recipient-agent", &test_signer()),
598            Err(TransportError::Capacity)
599        ));
600        assert_eq!(transport.inbox_depth("recipient-agent"), MAX_INBOX_SIZE);
601    }
602}
lean_ctx/core/capsule_transport.rs

lean_ctx/core/
capsule_transport.rs
