Skip to main content

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

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