Skip to main content

systemprompt_agent/services/a2a_server/processing/task_builder/
mod.rs

1mod builders;
2pub mod helpers;
3
4pub use builders::{
5    build_canceled_task, build_completed_task, build_mock_task, build_multiturn_task,
6    build_submitted_task,
7};
8
9use crate::models::a2a::{
10    Artifact, Message, MessageRole, Part, Task, TaskState, TaskStatus, TextPart,
11};
12use systemprompt_identifiers::{ContextId, MessageId, TaskId};
13use systemprompt_models::a2a::TaskMetadata;
14
15#[derive(Debug)]
16pub struct TaskBuilder {
17    task_id: TaskId,
18    context_id: ContextId,
19    state: TaskState,
20    response_text: String,
21    id: MessageId,
22    user_message: Option<Message>,
23    artifacts: Vec<Artifact>,
24    metadata: Option<TaskMetadata>,
25}
26
27impl TaskBuilder {
28    pub fn new(context_id: ContextId) -> Self {
29        Self {
30            task_id: TaskId::generate(),
31            context_id,
32            state: TaskState::Completed,
33            response_text: String::new(),
34            id: MessageId::generate(),
35            user_message: None,
36            artifacts: Vec::new(),
37            metadata: None,
38        }
39    }
40
41    pub fn with_task_id(mut self, task_id: TaskId) -> Self {
42        self.task_id = task_id;
43        self
44    }
45
46    pub const fn with_state(mut self, state: TaskState) -> Self {
47        self.state = state;
48        self
49    }
50
51    pub fn with_response_text(mut self, text: String) -> Self {
52        self.response_text = text;
53        self
54    }
55
56    pub fn with_message_id(mut self, id: impl Into<MessageId>) -> Self {
57        self.id = id.into();
58        self
59    }
60
61    pub fn with_user_message(mut self, message: Message) -> Self {
62        self.user_message = Some(message);
63        self
64    }
65
66    pub fn with_artifacts(mut self, artifacts: Vec<Artifact>) -> Self {
67        self.artifacts = artifacts;
68        self
69    }
70
71    pub fn with_metadata(mut self, metadata: TaskMetadata) -> Self {
72        self.metadata = Some(metadata);
73        self
74    }
75
76    pub fn build(self) -> Task {
77        let agent_message = Message {
78            role: MessageRole::Agent,
79            parts: vec![Part::Text(TextPart {
80                text: self.response_text.clone(),
81            })],
82            message_id: self.id.clone(),
83            task_id: Some(self.task_id.clone()),
84            context_id: self.context_id.clone(),
85            metadata: None,
86            extensions: None,
87            reference_task_ids: None,
88        };
89
90        let history = if let Some(user_msg) = self.user_message {
91            Some(vec![
92                user_msg,
93                Message {
94                    role: MessageRole::Agent,
95                    parts: vec![Part::Text(TextPart {
96                        text: self.response_text.clone(),
97                    })],
98                    message_id: MessageId::generate(),
99                    task_id: Some(self.task_id.clone()),
100                    context_id: self.context_id.clone(),
101                    metadata: None,
102                    extensions: None,
103                    reference_task_ids: None,
104                },
105            ])
106        } else {
107            None
108        };
109
110        Task {
111            id: self.task_id.clone(),
112            context_id: self.context_id.clone(),
113            status: TaskStatus {
114                state: self.state,
115                message: Some(agent_message),
116                timestamp: Some(chrono::Utc::now()),
117            },
118            history,
119            artifacts: if self.artifacts.is_empty() {
120                None
121            } else {
122                Some(self.artifacts)
123            },
124            metadata: self.metadata,
125            created_at: Some(chrono::Utc::now()),
126            last_modified: Some(chrono::Utc::now()),
127        }
128    }
129}