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 ft = file_object_to_filetype(f.file);
81                    if ft.mime_type().starts_with("image/") {
82                        parts.push(crate::Part::Image(ft));
83                    } else {
84                        parts.push(crate::Part::File(ft));
85                    }
86                }
87            }
88        }
89
90        let is_tool = parts.iter().any(|part| {
91            if let crate::Part::ToolResult(_) = part {
92                return true;
93            }
94            false
95        });
96
97        // Extract parts_metadata from message metadata if present
98        let parts_metadata: Option<crate::PartsMetadata> = message
99            .metadata
100            .as_ref()
101            .and_then(|m| m.get("parts"))
102            .and_then(|p| serde_json::from_value(p.clone()).ok());
103
104        Ok(crate::Message {
105            id: message.message_id.clone(),
106            role: if is_tool {
107                crate::MessageRole::Tool
108            } else {
109                match message.role {
110                    Role::User => crate::MessageRole::User,
111                    Role::Agent => crate::MessageRole::Assistant,
112                }
113            },
114            name: None,
115            parts,
116            parts_metadata,
117            ..Default::default()
118        })
119    }
120}
121
122impl TryFrom<distri_a2a::TaskStatusUpdateEvent> for crate::TaskEvent {
123    type Error = AgentError;
124
125    fn try_from(event: distri_a2a::TaskStatusUpdateEvent) -> Result<Self, Self::Error> {
126        let agent_event: crate::events::AgentEventType = event
127            .metadata
128            .ok_or_else(|| AgentError::Validation("missing metadata on status update".into()))
129            .and_then(|m| {
130                serde_json::from_value(m)
131                    .map_err(|e| AgentError::Validation(format!("invalid event metadata: {}", e)))
132            })?;
133
134        let created_at = event
135            .status
136            .timestamp
137            .and_then(|t| t.parse::<i64>().ok())
138            .unwrap_or(0);
139
140        Ok(crate::TaskEvent {
141            event: agent_event,
142            created_at,
143            is_final: event.r#final,
144        })
145    }
146}
147
148impl TryFrom<distri_a2a::MessageKind> for crate::TaskMessage {
149    type Error = AgentError;
150
151    fn try_from(mk: distri_a2a::MessageKind) -> Result<Self, Self::Error> {
152        match mk {
153            distri_a2a::MessageKind::Message(msg) => {
154                Ok(crate::TaskMessage::Message(crate::Message::try_from(msg)?))
155            }
156            distri_a2a::MessageKind::TaskStatusUpdate(evt) => {
157                Ok(crate::TaskMessage::Event(crate::TaskEvent::try_from(evt)?))
158            }
159            distri_a2a::MessageKind::Artifact(_) => Err(AgentError::Validation(
160                "artifact conversion not supported".into(),
161            )),
162        }
163    }
164}
165
166impl From<crate::TaskStatus> for TaskState {
167    fn from(status: crate::TaskStatus) -> Self {
168        match status {
169            crate::TaskStatus::Pending => TaskState::Submitted,
170            crate::TaskStatus::Running => TaskState::Working,
171            crate::TaskStatus::InputRequired => TaskState::InputRequired,
172            crate::TaskStatus::Completed => TaskState::Completed,
173            crate::TaskStatus::Failed => TaskState::Failed,
174            crate::TaskStatus::Canceled => TaskState::Canceled,
175        }
176    }
177}
178
179/// Convenience free-function wrapper around `From<crate::TaskStatus> for TaskState`.
180/// Intended for call sites that prefer explicit function-call syntax (e.g.
181/// `A2AService::prepare_resubscribe`, `orchestrator::get_thread`).
182pub fn map_task_status_to_a2a_state(status: &crate::TaskStatus) -> TaskState {
183    status.clone().into()
184}
185
186impl From<crate::Part> for Part {
187    fn from(part: crate::Part) -> Self {
188        match part {
189            crate::Part::Text(text) => Part::Text(TextPart { text }),
190            crate::Part::Image(image) => Part::File(FilePart {
191                file: filetype_to_fileobject(image),
192                metadata: None,
193            }),
194            crate::Part::File(file) => Part::File(FilePart {
195                file: filetype_to_fileobject(file),
196                metadata: None,
197            }),
198            // handle all the additional parts with a part_type
199            x => Part::Data(DataPart {
200                data: serde_json::to_value(x).unwrap(),
201            }),
202        }
203    }
204}
205
206fn file_object_to_filetype(file: FileObject) -> FileType {
207    match file {
208        FileObject::WithBytes {
209            bytes,
210            mime_type,
211            name,
212        } => FileType::Bytes {
213            bytes,
214            mime_type: mime_type.unwrap_or_default(),
215            name,
216        },
217        FileObject::WithUri {
218            uri,
219            mime_type,
220            name,
221        } => FileType::Url {
222            url: uri,
223            mime_type: mime_type.unwrap_or_default(),
224            name,
225        },
226    }
227}
228
229fn filetype_to_fileobject(file: FileType) -> FileObject {
230    match file {
231        FileType::Bytes {
232            bytes,
233            mime_type,
234            name,
235        } => FileObject::WithBytes {
236            bytes,
237            mime_type: Some(mime_type),
238            name: name.clone(),
239        },
240        FileType::Url {
241            url,
242            mime_type,
243            name,
244        } => FileObject::WithUri {
245            uri: url.clone(),
246            mime_type: Some(mime_type),
247            name: name.clone(),
248        },
249    }
250}
251
252impl From<crate::Task> for Task {
253    fn from(task: crate::Task) -> Self {
254        let history = vec![];
255        Task {
256            id: task.id.clone(),
257            status: TaskStatus {
258                state: match task.status {
259                    crate::TaskStatus::Pending => TaskState::Submitted,
260                    crate::TaskStatus::Running => TaskState::Working,
261                    crate::TaskStatus::InputRequired => TaskState::InputRequired,
262                    crate::TaskStatus::Completed => TaskState::Completed,
263                    crate::TaskStatus::Failed => TaskState::Failed,
264                    crate::TaskStatus::Canceled => TaskState::Canceled,
265                },
266                message: None,
267                timestamp: None,
268            },
269            kind: EventKind::Task,
270            context_id: task.thread_id.clone(),
271            artifacts: vec![],
272            history,
273            metadata: None,
274        }
275    }
276}
277
278impl From<crate::MessageRole> for Role {
279    fn from(role: crate::MessageRole) -> Self {
280        match role {
281            crate::MessageRole::User => Role::User,
282            crate::MessageRole::Assistant => Role::Agent,
283            // Developer messages are mapped to User for A2A protocol
284            // since they contain context that should be treated like user input
285            crate::MessageRole::Developer => Role::User,
286            _ => Role::Agent,
287        }
288    }
289}