Skip to main content

distri_types/
a2a_converters.rs

1use distri_a2a::{
2    DataPart, EventKind, FileObject, FilePart, Message, Part, Role, Task, TaskState, TaskStatus,
3    TextPart,
4};
5
6use serde::{Deserialize, Serialize};
7use serde_json::json;
8
9use crate::{AgentError, core::FileType};
10
11#[derive(Debug, Serialize, Deserialize, Clone)]
12#[serde(rename_all = "snake_case")]
13pub enum MessageMetadata {
14    Text,
15    Plan,
16    ToolCall,
17    ToolResult,
18}
19
20/// A2A Extension for agent metadata
21/// This allows tracking which agent generated each message
22#[derive(Debug, Serialize, Deserialize, Clone)]
23pub struct AgentMetadata {
24    /// The ID of the agent that generated this message
25    pub agent_id: String,
26    /// Optional agent name for display purposes
27    #[serde(skip_serializing_if = "Option::is_none")]
28    pub agent_name: Option<String>,
29}
30
31impl From<crate::Message> for MessageMetadata {
32    fn from(message: crate::Message) -> Self {
33        for part in message.parts.iter() {
34            match part {
35                crate::Part::ToolCall(_) => return MessageMetadata::ToolCall,
36                crate::Part::ToolResult(_) => return MessageMetadata::ToolResult,
37                _ => continue,
38            }
39        }
40        MessageMetadata::Text
41    }
42}
43
44impl TryFrom<Message> for crate::Message {
45    type Error = AgentError;
46
47    fn try_from(message: Message) -> Result<Self, Self::Error> {
48        let mut parts = Vec::new();
49        for part in message.parts {
50            match part {
51                Part::Text(t) => parts.push(crate::Part::Text(t.text.clone())),
52                Part::Data(d) => {
53                    if let Some(part_type) = d.data.get("part_type").and_then(|v| v.as_str()) {
54                        if let Some(data_content) = d.data.get("data") {
55                            // Create the properly structured object for
56
57                            let structured = json!({
58                                "part_type": part_type,
59                                "data": data_content
60                            });
61
62                            let part: crate::Part = serde_json::from_value(structured)?;
63                            parts.push(part);
64                        } else {
65                            return Err(AgentError::Validation(
66                                "Missing data
67                field for typed part"
68                                    .to_string(),
69                            ));
70                        }
71                    } else {
72                        return Err(AgentError::Validation(
73                            "Invalid part
74                type"
75                                .to_string(),
76                        ));
77                    }
78                }
79                Part::File(f) => {
80                    let mime_type = f.mime_type();
81                    if let Some(mime_type) = mime_type {
82                        if mime_type.starts_with("image/") {
83                            let ft = file_object_to_filetype(f.file.clone());
84                            parts.push(crate::Part::Image(ft));
85                        } else {
86                            return Err(AgentError::UnsupportedFileType(mime_type.to_string()));
87                        }
88                    } else {
89                        return Err(AgentError::UnsupportedFileType("unknown".to_string()));
90                    }
91                }
92            }
93        }
94
95        let is_tool = parts.iter().any(|part| {
96            if let crate::Part::ToolResult(_) = part {
97                return true;
98            }
99            false
100        });
101
102        // Extract parts_metadata from message metadata if present
103        let parts_metadata: Option<crate::PartsMetadata> = message
104            .metadata
105            .as_ref()
106            .and_then(|m| m.get("parts"))
107            .and_then(|p| serde_json::from_value(p.clone()).ok());
108
109        Ok(crate::Message {
110            id: message.message_id.clone(),
111            role: if is_tool {
112                crate::MessageRole::Tool
113            } else {
114                match message.role {
115                    Role::User => crate::MessageRole::User,
116                    Role::Agent => crate::MessageRole::Assistant,
117                }
118            },
119            name: None,
120            parts,
121            parts_metadata,
122            ..Default::default()
123        })
124    }
125}
126
127impl TryFrom<distri_a2a::TaskStatusUpdateEvent> for crate::TaskEvent {
128    type Error = AgentError;
129
130    fn try_from(event: distri_a2a::TaskStatusUpdateEvent) -> Result<Self, Self::Error> {
131        let agent_event: crate::events::AgentEventType = event
132            .metadata
133            .ok_or_else(|| AgentError::Validation("missing metadata on status update".into()))
134            .and_then(|m| {
135                serde_json::from_value(m)
136                    .map_err(|e| AgentError::Validation(format!("invalid event metadata: {}", e)))
137            })?;
138
139        let created_at = event
140            .status
141            .timestamp
142            .and_then(|t| t.parse::<i64>().ok())
143            .unwrap_or(0);
144
145        Ok(crate::TaskEvent {
146            event: agent_event,
147            created_at,
148            is_final: event.r#final,
149        })
150    }
151}
152
153impl TryFrom<distri_a2a::MessageKind> for crate::TaskMessage {
154    type Error = AgentError;
155
156    fn try_from(mk: distri_a2a::MessageKind) -> Result<Self, Self::Error> {
157        match mk {
158            distri_a2a::MessageKind::Message(msg) => {
159                Ok(crate::TaskMessage::Message(crate::Message::try_from(msg)?))
160            }
161            distri_a2a::MessageKind::TaskStatusUpdate(evt) => {
162                Ok(crate::TaskMessage::Event(crate::TaskEvent::try_from(evt)?))
163            }
164            distri_a2a::MessageKind::Artifact(_) => Err(AgentError::Validation(
165                "artifact conversion not supported".into(),
166            )),
167        }
168    }
169}
170
171impl From<crate::TaskStatus> for TaskState {
172    fn from(status: crate::TaskStatus) -> Self {
173        match status {
174            crate::TaskStatus::Pending => TaskState::Submitted,
175            crate::TaskStatus::Running => TaskState::Working,
176            crate::TaskStatus::InputRequired => TaskState::InputRequired,
177            crate::TaskStatus::Completed => TaskState::Completed,
178            crate::TaskStatus::Failed => TaskState::Failed,
179            crate::TaskStatus::Canceled => TaskState::Canceled,
180        }
181    }
182}
183
184impl From<crate::Part> for Part {
185    fn from(part: crate::Part) -> Self {
186        match part {
187            crate::Part::Text(text) => Part::Text(TextPart { text }),
188            crate::Part::Image(image) => Part::File(FilePart {
189                file: filetype_to_fileobject(image),
190                metadata: None,
191            }),
192
193            // handle all  the additional parts with a part_type
194            x => Part::Data(DataPart {
195                data: serde_json::to_value(x).unwrap(),
196            }),
197        }
198    }
199}
200
201fn file_object_to_filetype(file: FileObject) -> FileType {
202    match file {
203        FileObject::WithBytes {
204            bytes,
205            mime_type,
206            name,
207        } => FileType::Bytes {
208            bytes,
209            mime_type: mime_type.unwrap_or_default(),
210            name,
211        },
212        FileObject::WithUri {
213            uri,
214            mime_type,
215            name,
216        } => FileType::Url {
217            url: uri,
218            mime_type: mime_type.unwrap_or_default(),
219            name,
220        },
221    }
222}
223
224fn filetype_to_fileobject(file: FileType) -> FileObject {
225    match file {
226        FileType::Bytes {
227            bytes,
228            mime_type,
229            name,
230        } => FileObject::WithBytes {
231            bytes,
232            mime_type: Some(mime_type),
233            name: name.clone(),
234        },
235        FileType::Url {
236            url,
237            mime_type,
238            name,
239        } => FileObject::WithUri {
240            uri: url.clone(),
241            mime_type: Some(mime_type),
242            name: name.clone(),
243        },
244    }
245}
246
247impl From<crate::Task> for Task {
248    fn from(task: crate::Task) -> Self {
249        let history = vec![];
250        Task {
251            id: task.id.clone(),
252            status: TaskStatus {
253                state: match task.status {
254                    crate::TaskStatus::Pending => TaskState::Submitted,
255                    crate::TaskStatus::Running => TaskState::Working,
256                    crate::TaskStatus::InputRequired => TaskState::InputRequired,
257                    crate::TaskStatus::Completed => TaskState::Completed,
258                    crate::TaskStatus::Failed => TaskState::Failed,
259                    crate::TaskStatus::Canceled => TaskState::Canceled,
260                },
261                message: None,
262                timestamp: None,
263            },
264            kind: EventKind::Task,
265            context_id: task.thread_id.clone(),
266            artifacts: vec![],
267            history,
268            metadata: None,
269        }
270    }
271}
272
273impl From<crate::MessageRole> for Role {
274    fn from(role: crate::MessageRole) -> Self {
275        match role {
276            crate::MessageRole::User => Role::User,
277            crate::MessageRole::Assistant => Role::Agent,
278            // Developer messages are mapped to User for A2A protocol
279            // since they contain context that should be treated like user input
280            crate::MessageRole::Developer => Role::User,
281            _ => Role::Agent,
282        }
283    }
284}