1use chrono::{DateTime, Utc};
9use hashbrown::HashMap;
10use serde::{Deserialize, Serialize};
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
18#[serde(rename_all = "kebab-case")]
19pub enum TaskState {
20 #[default]
22 Submitted,
23 Working,
25 InputRequired,
27 Completed,
29 Failed,
31 Canceled,
33 Rejected,
35 AuthRequired,
37 #[serde(other)]
39 Unknown,
40}
41
42impl TaskState {
43 fn is_terminal(&self) -> bool {
45 matches!(self, TaskState::Completed | TaskState::Failed | TaskState::Canceled | TaskState::Rejected)
46 }
47
48 fn is_cancelable(&self) -> bool {
50 matches!(self, TaskState::Submitted | TaskState::Working | TaskState::InputRequired)
51 }
52}
53
54#[derive(Debug, Clone, Serialize, Deserialize)]
56pub struct TaskStatus {
57 pub state: TaskState,
59 #[serde(skip_serializing_if = "Option::is_none")]
61 pub message: Option<Message>,
62 pub(crate) timestamp: DateTime<Utc>,
64}
65
66impl TaskStatus {
67 pub(crate) fn new(state: TaskState) -> Self {
69 Self { state, message: None, timestamp: Utc::now() }
70 }
71
72 pub(crate) fn with_message(state: TaskState, message: Message) -> Self {
74 Self {
75 state,
76 message: Some(message),
77 timestamp: Utc::now(),
78 }
79 }
80}
81
82impl Default for TaskStatus {
83 fn default() -> Self {
84 Self::new(TaskState::Submitted)
85 }
86}
87
88#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
94#[serde(rename_all = "lowercase")]
95pub enum MessageRole {
96 User,
98 Agent,
100}
101
102#[derive(Debug, Clone, Serialize, Deserialize)]
104#[serde(rename_all = "camelCase")]
105pub struct Message {
106 pub role: MessageRole,
108 pub parts: Vec<Part>,
110 #[serde(skip_serializing_if = "Option::is_none")]
112 message_id: Option<String>,
113 #[serde(skip_serializing_if = "Option::is_none")]
115 task_id: Option<String>,
116 #[serde(skip_serializing_if = "Option::is_none")]
118 context_id: Option<String>,
119 #[serde(skip_serializing_if = "Vec::is_empty", default)]
121 reference_task_ids: Vec<String>,
122 #[serde(skip_serializing_if = "Option::is_none")]
124 metadata: Option<HashMap<String, serde_json::Value>>,
125}
126
127impl Message {
128 fn new(role: MessageRole, parts: Vec<Part>) -> Self {
130 Self {
131 role,
132 parts,
133 message_id: None,
134 task_id: None,
135 context_id: None,
136 reference_task_ids: Vec::new(),
137 metadata: None,
138 }
139 }
140
141 pub(crate) fn agent_text(text: impl Into<String>) -> Self {
143 Self::new(MessageRole::Agent, vec![Part::text(text)])
144 }
145
146 pub fn user_text(text: impl Into<String>) -> Self {
148 Self::new(MessageRole::User, vec![Part::text(text)])
149 }
150
151 pub fn with_id(mut self, id: impl Into<String>) -> Self {
153 self.message_id = Some(id.into());
154 self
155 }
156
157 pub fn with_task_id(mut self, task_id: impl Into<String>) -> Self {
159 self.task_id = Some(task_id.into());
160 self
161 }
162
163 pub fn with_context_id(mut self, context_id: impl Into<String>) -> Self {
165 self.context_id = Some(context_id.into());
166 self
167 }
168}
169
170#[derive(Debug, Clone, Serialize, Deserialize)]
172#[serde(tag = "type", rename_all = "camelCase")]
173pub enum Part {
174 #[serde(rename = "text")]
176 Text {
177 text: String,
179 },
180 #[serde(rename = "file")]
182 File {
183 file: FileContent,
185 },
186 #[serde(rename = "data")]
188 Data {
189 data: serde_json::Value,
191 },
192 #[serde(other)]
194 Unknown,
195}
196
197impl Part {
198 fn text(text: impl Into<String>) -> Self {
200 Part::Text { text: text.into() }
201 }
202
203 pub fn file_uri(uri: impl Into<String>, mime_type: Option<String>) -> Self {
205 Part::File {
206 file: FileContent::Uri { uri: uri.into(), mime_type },
207 }
208 }
209
210 pub fn file_bytes(bytes: Vec<u8>, mime_type: Option<String>, name: Option<String>) -> Self {
212 Part::File {
213 file: FileContent::Bytes {
214 bytes: base64::Engine::encode(&base64::engine::general_purpose::STANDARD, &bytes),
215 mime_type,
216 name,
217 },
218 }
219 }
220
221 fn data(data: serde_json::Value) -> Self {
223 Part::Data { data }
224 }
225
226 pub fn as_text(&self) -> Option<&str> {
228 match self {
229 Part::Text { text } => Some(text),
230 _ => None,
231 }
232 }
233
234 pub fn is_unknown(&self) -> bool {
236 matches!(self, Part::Unknown)
237 }
238}
239
240#[derive(Debug, Clone, Serialize, Deserialize)]
242#[serde(untagged)]
243pub enum FileContent {
244 Uri {
246 uri: String,
248 #[serde(skip_serializing_if = "Option::is_none")]
250 mime_type: Option<String>,
251 },
252 Bytes {
254 bytes: String,
256 #[serde(skip_serializing_if = "Option::is_none")]
258 mime_type: Option<String>,
259 #[serde(skip_serializing_if = "Option::is_none")]
261 name: Option<String>,
262 },
263}
264
265#[derive(Debug, Clone, Serialize, Deserialize)]
271#[serde(rename_all = "camelCase")]
272pub struct Artifact {
273 pub id: String,
275 #[serde(skip_serializing_if = "Option::is_none")]
277 name: Option<String>,
278 #[serde(skip_serializing_if = "Option::is_none")]
280 description: Option<String>,
281 pub parts: Vec<Part>,
283 #[serde(skip_serializing_if = "Option::is_none")]
285 index: Option<u32>,
286 #[serde(skip_serializing_if = "Option::is_none")]
288 metadata: Option<HashMap<String, serde_json::Value>>,
289}
290
291impl Artifact {
292 pub(crate) fn text(id: impl Into<String>, text: impl Into<String>) -> Self {
294 Self {
295 id: id.into(),
296 name: None,
297 description: None,
298 parts: vec![Part::text(text)],
299 index: None,
300 metadata: None,
301 }
302 }
303
304 pub fn file(id: impl Into<String>, file: FileContent) -> Self {
306 Self {
307 id: id.into(),
308 name: None,
309 description: None,
310 parts: vec![Part::File { file }],
311 index: None,
312 metadata: None,
313 }
314 }
315
316 fn with_name(mut self, name: impl Into<String>) -> Self {
318 self.name = Some(name.into());
319 self
320 }
321
322 fn with_description(mut self, description: impl Into<String>) -> Self {
324 self.description = Some(description.into());
325 self
326 }
327}
328
329#[derive(Debug, Clone, Serialize, Deserialize)]
335#[serde(rename_all = "camelCase")]
336pub struct Task {
337 pub id: String,
339 #[serde(skip_serializing_if = "Option::is_none")]
341 pub context_id: Option<String>,
342 pub status: TaskStatus,
344 #[serde(skip_serializing_if = "Vec::is_empty", default)]
346 pub artifacts: Vec<Artifact>,
347 #[serde(skip_serializing_if = "Vec::is_empty", default)]
349 pub history: Vec<Message>,
350 #[serde(skip_serializing_if = "Option::is_none")]
352 metadata: Option<HashMap<String, serde_json::Value>>,
353 #[serde(default = "default_task_kind")]
355 kind: String,
356}
357
358fn default_task_kind() -> String {
359 "task".to_string()
360}
361
362impl Task {
363 pub(crate) fn new() -> Self {
365 Self {
366 id: uuid::Uuid::new_v4().to_string(),
367 context_id: None,
368 status: TaskStatus::default(),
369 artifacts: Vec::new(),
370 history: Vec::new(),
371 metadata: None,
372 kind: "task".to_string(),
373 }
374 }
375
376 pub fn with_id(id: impl Into<String>) -> Self {
378 Self {
379 id: id.into(),
380 context_id: None,
381 status: TaskStatus::default(),
382 artifacts: Vec::new(),
383 history: Vec::new(),
384 metadata: None,
385 kind: "task".to_string(),
386 }
387 }
388
389 pub(crate) fn with_context_id(mut self, context_id: impl Into<String>) -> Self {
391 self.context_id = Some(context_id.into());
392 self
393 }
394
395 pub(crate) fn state(&self) -> TaskState {
397 self.status.state
398 }
399
400 pub(crate) fn is_terminal(&self) -> bool {
402 self.status.state.is_terminal()
403 }
404
405 pub(crate) fn is_cancelable(&self) -> bool {
407 self.status.state.is_cancelable()
408 }
409
410 fn update_status(&mut self, state: TaskState, message: Option<Message>) {
412 self.status = match message {
413 Some(msg) => TaskStatus::with_message(state, msg),
414 None => TaskStatus::new(state),
415 };
416 }
417
418 pub fn add_artifact(&mut self, artifact: Artifact) {
420 self.artifacts.push(artifact);
421 }
422
423 pub fn add_message(&mut self, message: Message) {
425 self.history.push(message);
426 }
427}
428
429impl Default for Task {
430 fn default() -> Self {
431 Self::new()
432 }
433}
434
435#[cfg(test)]
436mod tests {
437 use super::*;
438
439 #[test]
440 fn test_task_state_terminal() {
441 assert!(!TaskState::Submitted.is_terminal());
442 assert!(!TaskState::Working.is_terminal());
443 assert!(TaskState::Completed.is_terminal());
444 assert!(TaskState::Failed.is_terminal());
445 assert!(TaskState::Canceled.is_terminal());
446 }
447
448 #[test]
449 fn test_task_state_cancelable() {
450 assert!(TaskState::Submitted.is_cancelable());
451 assert!(TaskState::Working.is_cancelable());
452 assert!(TaskState::InputRequired.is_cancelable());
453 assert!(!TaskState::Completed.is_cancelable());
454 assert!(!TaskState::Failed.is_cancelable());
455 }
456
457 #[test]
458 fn test_message_creation() {
459 let msg = Message::agent_text("Hello, world!");
460 assert_eq!(msg.role, MessageRole::Agent);
461 assert_eq!(msg.parts.len(), 1);
462 assert_eq!(msg.parts[0].as_text(), Some("Hello, world!"));
463 }
464
465 #[test]
466 fn test_task_lifecycle() {
467 let mut task = Task::new();
468 assert_eq!(task.state(), TaskState::Submitted);
469 assert!(!task.is_terminal());
470 assert!(task.is_cancelable());
471
472 task.update_status(TaskState::Working, None);
473 assert_eq!(task.state(), TaskState::Working);
474
475 task.update_status(TaskState::Completed, Some(Message::agent_text("Task completed")));
476 assert!(task.is_terminal());
477 assert!(!task.is_cancelable());
478 }
479
480 #[test]
481 fn test_part_serialization() {
482 let text_part = Part::text("Hello");
483 let json = serde_json::to_string(&text_part).expect("serialize");
484 assert!(json.contains("\"type\":\"text\""));
485
486 let data_part = Part::data(serde_json::json!({"key": "value"}));
487 let json = serde_json::to_string(&data_part).expect("serialize");
488 assert!(json.contains("\"type\":\"data\""));
489 }
490
491 #[test]
492 fn test_artifact_creation() {
493 let artifact = Artifact::text("art-1", "Generated content")
494 .with_name("output.txt")
495 .with_description("The generated output file");
496
497 assert_eq!(artifact.id, "art-1");
498 assert_eq!(artifact.name, Some("output.txt".to_string()));
499 assert_eq!(artifact.parts.len(), 1);
500 }
501}