Skip to main content

lago_core/
protocol_bridge.rs

1//! Bridge between Lago's internal types and the canonical `aios-protocol` types.
2//!
3//! Provides `From`/`Into` conversions for IDs and event payloads so that
4//! Lago can interoperate with the Agent OS protocol without changing its
5//! internal storage format.
6
7use crate::event::{EventEnvelope, EventPayload};
8use crate::id::*;
9
10/// Re-export the canonical protocol for downstream convenience.
11pub use aios_protocol;
12
13// ─── ID conversions: Lago (ULID String) ↔ Protocol (String) ────────────────
14//
15// Both use String-based IDs, so conversions are trivial.
16
17impl From<SessionId> for aios_protocol::SessionId {
18    fn from(id: SessionId) -> Self {
19        aios_protocol::SessionId::from_string(id.as_str())
20    }
21}
22
23impl From<aios_protocol::SessionId> for SessionId {
24    fn from(id: aios_protocol::SessionId) -> Self {
25        SessionId::from_string(id.as_str())
26    }
27}
28
29impl From<EventId> for aios_protocol::EventId {
30    fn from(id: EventId) -> Self {
31        aios_protocol::EventId::from_string(id.as_str())
32    }
33}
34
35impl From<aios_protocol::EventId> for EventId {
36    fn from(id: aios_protocol::EventId) -> Self {
37        EventId::from_string(id.as_str())
38    }
39}
40
41impl From<BranchId> for aios_protocol::BranchId {
42    fn from(id: BranchId) -> Self {
43        aios_protocol::BranchId::from_string(id.as_str())
44    }
45}
46
47impl From<aios_protocol::BranchId> for BranchId {
48    fn from(id: aios_protocol::BranchId) -> Self {
49        BranchId::from_string(id.as_str())
50    }
51}
52
53impl From<RunId> for aios_protocol::RunId {
54    fn from(id: RunId) -> Self {
55        aios_protocol::RunId::from_string(id.as_str())
56    }
57}
58
59impl From<aios_protocol::RunId> for RunId {
60    fn from(id: aios_protocol::RunId) -> Self {
61        RunId::from_string(id.as_str())
62    }
63}
64
65impl From<SnapshotId> for aios_protocol::SnapshotId {
66    fn from(id: SnapshotId) -> Self {
67        aios_protocol::SnapshotId::from_string(id.as_str())
68    }
69}
70
71impl From<ApprovalId> for aios_protocol::ApprovalId {
72    fn from(id: ApprovalId) -> Self {
73        aios_protocol::ApprovalId::from_string(id.as_str())
74    }
75}
76
77impl From<MemoryId> for aios_protocol::MemoryId {
78    fn from(id: MemoryId) -> Self {
79        aios_protocol::MemoryId::from_string(id.as_str())
80    }
81}
82
83impl From<BlobHash> for aios_protocol::BlobHash {
84    fn from(hash: BlobHash) -> Self {
85        aios_protocol::BlobHash::from_hex(hash.as_str())
86    }
87}
88
89impl From<aios_protocol::BlobHash> for BlobHash {
90    fn from(hash: aios_protocol::BlobHash) -> Self {
91        BlobHash::from_hex(hash.as_str())
92    }
93}
94
95/// Convert a Lago `EventEnvelope` to a canonical `aios_protocol::EventEnvelope`.
96///
97/// Uses JSON round-trip for the payload to handle all current and future
98/// event variants without maintaining a manual mapping.
99impl EventEnvelope {
100    pub fn to_protocol(&self) -> Option<aios_protocol::EventEnvelope> {
101        let kind_json = serde_json::to_value(&self.payload).ok()?;
102        let protocol_kind: aios_protocol::EventKind = serde_json::from_value(kind_json).ok()?;
103
104        Some(aios_protocol::EventEnvelope {
105            event_id: self.event_id.clone().into(),
106            session_id: self.session_id.clone().into(),
107            agent_id: self
108                .metadata
109                .get("agent_id")
110                .map(|s| aios_protocol::AgentId::from_string(s.clone()))
111                .unwrap_or_default(),
112            branch_id: self.branch_id.clone().into(),
113            run_id: self.run_id.clone().map(Into::into),
114            seq: self.seq,
115            timestamp: self.timestamp,
116            actor: self
117                .metadata
118                .get("actor")
119                .and_then(|v| serde_json::from_str::<aios_protocol::EventActor>(v).ok())
120                .unwrap_or_default(),
121            schema: self
122                .metadata
123                .get("schema")
124                .and_then(|v| serde_json::from_str::<aios_protocol::EventSchema>(v).ok())
125                .unwrap_or_default(),
126            parent_id: self.parent_id.clone().map(Into::into),
127            trace_id: self.metadata.get("trace_id").cloned(),
128            span_id: self.metadata.get("span_id").cloned(),
129            digest: self.metadata.get("digest").cloned(),
130            kind: protocol_kind,
131            metadata: self.metadata.clone(),
132            schema_version: self.schema_version,
133        })
134    }
135}
136
137/// Convert a canonical `aios_protocol::EventEnvelope` back to Lago's format.
138pub fn from_protocol(envelope: &aios_protocol::EventEnvelope) -> Option<EventEnvelope> {
139    let kind_json = serde_json::to_value(&envelope.kind).ok()?;
140    let lago_payload: EventPayload = serde_json::from_value(kind_json).ok()?;
141    let mut metadata = envelope.metadata.clone();
142    metadata
143        .entry("agent_id".to_string())
144        .or_insert_with(|| envelope.agent_id.to_string());
145    metadata
146        .entry("actor".to_string())
147        .or_insert_with(|| serde_json::to_string(&envelope.actor).unwrap_or_default());
148    metadata
149        .entry("schema".to_string())
150        .or_insert_with(|| serde_json::to_string(&envelope.schema).unwrap_or_default());
151    if let Some(trace_id) = &envelope.trace_id {
152        metadata.insert("trace_id".to_string(), trace_id.clone());
153    }
154    if let Some(span_id) = &envelope.span_id {
155        metadata.insert("span_id".to_string(), span_id.clone());
156    }
157    if let Some(digest) = &envelope.digest {
158        metadata.insert("digest".to_string(), digest.clone());
159    }
160
161    Some(EventEnvelope {
162        event_id: EventId::from_string(envelope.event_id.as_str()),
163        session_id: SessionId::from_string(envelope.session_id.as_str()),
164        branch_id: BranchId::from_string(envelope.branch_id.as_str()),
165        run_id: envelope
166            .run_id
167            .as_ref()
168            .map(|id| RunId::from_string(id.as_str())),
169        seq: envelope.seq,
170        timestamp: envelope.timestamp,
171        parent_id: envelope
172            .parent_id
173            .as_ref()
174            .map(|id| EventId::from_string(id.as_str())),
175        payload: lago_payload,
176        metadata,
177        schema_version: envelope.schema_version,
178    })
179}
180
181#[cfg(test)]
182mod tests {
183    use super::*;
184    use crate::event::EventPayload;
185    use std::collections::HashMap;
186
187    #[test]
188    fn session_id_roundtrip() {
189        let lago_id = SessionId::from_string("SESS001");
190        let proto_id: aios_protocol::SessionId = lago_id.clone().into();
191        assert_eq!(proto_id.as_str(), "SESS001");
192        let back: SessionId = proto_id.into();
193        assert_eq!(back, lago_id);
194    }
195
196    #[test]
197    fn blob_hash_roundtrip() {
198        let lago_hash = BlobHash::from_hex("deadbeef");
199        let proto_hash: aios_protocol::BlobHash = lago_hash.clone().into();
200        assert_eq!(proto_hash.as_str(), "deadbeef");
201        let back: BlobHash = proto_hash.into();
202        assert_eq!(back, lago_hash);
203    }
204
205    #[test]
206    fn envelope_to_protocol_roundtrip() {
207        let envelope = EventEnvelope {
208            event_id: EventId::from_string("EVT001"),
209            session_id: SessionId::from_string("SESS001"),
210            branch_id: BranchId::from_string("main"),
211            run_id: None,
212            seq: 42,
213            timestamp: 1_700_000_000_000_000,
214            parent_id: None,
215            payload: EventPayload::RunStarted {
216                provider: "anthropic".into(),
217                max_iterations: 10,
218            },
219            metadata: HashMap::new(),
220            schema_version: 1,
221        };
222
223        let proto = envelope.to_protocol().expect("convert to protocol");
224        assert_eq!(proto.seq, 42);
225        assert_eq!(proto.event_id.as_str(), "EVT001");
226        assert!(matches!(
227            proto.kind,
228            aios_protocol::EventKind::RunStarted { .. }
229        ));
230
231        // Convert back
232        let back = from_protocol(&proto).expect("convert from protocol");
233        assert_eq!(back.seq, 42);
234        assert!(matches!(back.payload, EventPayload::RunStarted { .. }));
235    }
236
237    #[test]
238    fn lago_error_raised_converts_to_protocol() {
239        let envelope = EventEnvelope {
240            event_id: EventId::from_string("EVT002"),
241            session_id: SessionId::from_string("SESS001"),
242            branch_id: BranchId::from_string("main"),
243            run_id: None,
244            seq: 1,
245            timestamp: 1_700_000_000_000_000,
246            parent_id: None,
247            payload: EventPayload::ErrorRaised {
248                message: "test".into(),
249            },
250            metadata: HashMap::new(),
251            schema_version: 1,
252        };
253
254        let proto = envelope.to_protocol().expect("convert to protocol");
255        assert!(matches!(
256            proto.kind,
257            aios_protocol::EventKind::ErrorRaised { .. }
258        ));
259    }
260}