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#[derive(Debug, Serialize, Deserialize, Clone)]
23pub struct AgentMetadata {
24 pub agent_id: String,
26 #[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 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 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
179pub 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 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 crate::MessageRole::Developer => Role::User,
286 _ => Role::Agent,
287 }
288 }
289}