systemprompt_agent/services/a2a_server/processing/task_builder/
builders.rs1use super::TaskBuilder;
2use super::helpers::{content_to_json, extract_text_from_content};
3use crate::models::a2a::{
4 Artifact, DataPart, Message, MessageRole, Part, Task, TaskState, TaskStatus, TextPart,
5};
6use crate::services::mcp::parse_tool_response;
7use serde_json::json;
8use systemprompt_identifiers::{ContextId, MessageId, TaskId};
9use systemprompt_models::a2a::{ArtifactMetadata, TaskMetadata, agent_names};
10use systemprompt_models::{CallToolResult, ToolCall};
11
12pub fn build_completed_task(
13 task_id: TaskId,
14 context_id: ContextId,
15 response_text: String,
16 user_message: Message,
17 artifacts: Vec<Artifact>,
18) -> Task {
19 TaskBuilder::new(context_id)
20 .with_task_id(task_id)
21 .with_state(TaskState::Completed)
22 .with_response_text(response_text)
23 .with_user_message(user_message)
24 .with_artifacts(artifacts)
25 .build()
26}
27
28pub fn build_canceled_task(task_id: TaskId, context_id: ContextId) -> Task {
29 TaskBuilder::new(context_id)
30 .with_task_id(task_id)
31 .with_state(TaskState::Canceled)
32 .with_response_text("Task was canceled.".to_string())
33 .build()
34}
35
36pub fn build_mock_task(task_id: TaskId) -> Task {
37 let mock_context_id = ContextId::generate();
38 TaskBuilder::new(mock_context_id)
39 .with_task_id(task_id)
40 .with_state(TaskState::Completed)
41 .with_response_text("Task completed successfully.".to_string())
42 .build()
43}
44
45pub fn build_submitted_task(
46 task_id: TaskId,
47 context_id: ContextId,
48 user_message: Message,
49 agent_name: &str,
50) -> Task {
51 Task {
52 id: task_id,
53 context_id,
54 status: TaskStatus {
55 state: TaskState::Submitted,
56 message: None,
57 timestamp: Some(chrono::Utc::now()),
58 },
59 history: Some(vec![user_message]),
60 artifacts: None,
61 metadata: Some(TaskMetadata::new_agent_message(agent_name.to_string())),
62 created_at: Some(chrono::Utc::now()),
63 last_modified: Some(chrono::Utc::now()),
64 }
65}
66
67#[derive(Debug)]
68pub struct BuildMultiturnTaskParams {
69 pub context_id: ContextId,
70 pub task_id: TaskId,
71 pub user_message: Message,
72 pub tool_calls: Vec<ToolCall>,
73 pub tool_results: Vec<CallToolResult>,
74 pub final_response: String,
75 pub total_iterations: usize,
76}
77
78pub fn build_multiturn_task(params: BuildMultiturnTaskParams) -> Task {
79 let BuildMultiturnTaskParams {
80 context_id,
81 task_id,
82 user_message,
83 tool_calls,
84 tool_results,
85 final_response,
86 total_iterations,
87 } = params;
88 let ctx_id = context_id;
89
90 let history = build_history(BuildHistoryParams {
91 ctx_id: &ctx_id,
92 task_id: &task_id,
93 user_message,
94 tool_calls: &tool_calls,
95 tool_results: &tool_results,
96 final_response: &final_response,
97 });
98
99 let artifacts = build_artifacts(&ctx_id, &task_id, &tool_calls, &tool_results);
100
101 Task {
102 id: task_id.clone(),
103 context_id: ctx_id.clone(),
104 status: TaskStatus {
105 state: TaskState::Completed,
106 message: Some(Message {
107 role: MessageRole::Agent,
108 parts: vec![Part::Text(TextPart {
109 text: final_response,
110 })],
111 message_id: MessageId::generate(),
112 task_id: Some(task_id),
113 context_id: ctx_id,
114 metadata: None,
115 extensions: None,
116 reference_task_ids: None,
117 }),
118 timestamp: Some(chrono::Utc::now()),
119 },
120 history: Some(history),
121 artifacts: if artifacts.is_empty() {
122 None
123 } else {
124 Some(artifacts)
125 },
126 metadata: Some(
127 TaskMetadata::new_agent_message(agent_names::SYSTEM.to_string())
128 .with_extension("total_iterations".to_string(), json!(total_iterations))
129 .with_extension("total_tools_called".to_string(), json!(tool_calls.len())),
130 ),
131 created_at: Some(chrono::Utc::now()),
132 last_modified: Some(chrono::Utc::now()),
133 }
134}
135
136struct BuildHistoryParams<'a> {
137 ctx_id: &'a ContextId,
138 task_id: &'a TaskId,
139 user_message: Message,
140 tool_calls: &'a [ToolCall],
141 tool_results: &'a [CallToolResult],
142 final_response: &'a str,
143}
144
145fn build_history(params: BuildHistoryParams<'_>) -> Vec<Message> {
146 let BuildHistoryParams {
147 ctx_id,
148 task_id,
149 user_message,
150 tool_calls,
151 tool_results,
152 final_response,
153 } = params;
154 let mut history = Vec::new();
155 history.push(user_message);
156
157 let mut iteration = 1;
158 let mut call_idx = 0;
159
160 while call_idx < tool_calls.len() {
161 let iteration_calls: Vec<_> = tool_calls
162 .iter()
163 .skip(call_idx)
164 .take_while(|_| call_idx < tool_calls.len())
165 .cloned()
166 .collect();
167
168 if iteration_calls.is_empty() {
169 break;
170 }
171
172 history.push(Message {
173 role: MessageRole::Agent,
174 parts: vec![Part::Text(TextPart {
175 text: format!("Executing {} tool(s)...", iteration_calls.len()),
176 })],
177 message_id: MessageId::generate(),
178 task_id: Some(task_id.clone()),
179 context_id: ctx_id.clone(),
180 metadata: Some(json!({
181 "iteration": iteration,
182 "tool_calls": iteration_calls.iter().map(|tc| {
183 json!({"id": tc.ai_tool_call_id.as_ref(), "name": tc.name})
184 }).collect::<Vec<_>>()
185 })),
186 extensions: None,
187 reference_task_ids: None,
188 });
189
190 let results_text = iteration_calls
191 .iter()
192 .enumerate()
193 .filter_map(|(idx, call)| {
194 let result_idx = call_idx + idx;
195 tool_results.get(result_idx).map(|r| {
196 let content_text = extract_text_from_content(&r.content);
197 format!("Tool '{}' result: {}", call.name, content_text)
198 })
199 })
200 .collect::<Vec<_>>()
201 .join("\n");
202
203 history.push(Message {
204 role: MessageRole::User,
205 parts: vec![Part::Text(TextPart { text: results_text })],
206 message_id: MessageId::generate(),
207 task_id: Some(task_id.clone()),
208 context_id: ctx_id.clone(),
209 metadata: Some(json!({
210 "iteration": iteration,
211 "tool_results": true
212 })),
213 extensions: None,
214 reference_task_ids: None,
215 });
216
217 call_idx += iteration_calls.len();
218 iteration += 1;
219 }
220
221 history.push(Message {
222 role: MessageRole::Agent,
223 parts: vec![Part::Text(TextPart {
224 text: final_response.to_string(),
225 })],
226 message_id: MessageId::generate(),
227 task_id: Some(task_id.clone()),
228 context_id: ctx_id.clone(),
229 metadata: Some(json!({
230 "iteration": iteration,
231 "final_synthesis": true
232 })),
233 extensions: None,
234 reference_task_ids: None,
235 });
236
237 history
238}
239
240fn build_artifacts(
241 ctx_id: &ContextId,
242 task_id: &TaskId,
243 tool_calls: &[ToolCall],
244 tool_results: &[CallToolResult],
245) -> Vec<Artifact> {
246 tool_results
247 .iter()
248 .enumerate()
249 .filter_map(|(idx, result)| {
250 let tool_call = tool_calls.get(idx)?;
251 let tool_name = &tool_call.name;
252 let call_id = tool_call.ai_tool_call_id.as_ref();
253 let is_error = result.is_error?;
254
255 let structured_content = result.structured_content.as_ref()?;
256 let parsed = parse_tool_response(structured_content)
257 .map_err(|e| {
258 tracing::debug!(tool_name = %tool_name, error = %e, "Failed to parse tool response, skipping artifact");
259 e
260 })
261 .ok()?;
262
263 let mut data_map = serde_json::Map::new();
264 data_map.insert("call_id".to_string(), json!(call_id));
265 data_map.insert("tool_name".to_string(), json!(tool_name));
266 data_map.insert("output".to_string(), content_to_json(&result.content));
267 data_map.insert(
268 "status".to_string(),
269 json!(if is_error { "error" } else { "success" }),
270 );
271
272 Some(Artifact {
273 id: parsed.artifact_id,
274 title: Some(format!("tool_execution_{}", idx + 1)),
275 description: Some(format!("Result from tool: {tool_name}")),
276 parts: vec![Part::Data(DataPart { data: data_map })],
277 extensions: vec![],
278 metadata: ArtifactMetadata::new(
279 "tool_execution".to_string(),
280 ctx_id.clone(),
281 task_id.clone(),
282 )
283 .with_mcp_execution_id(call_id.to_string())
284 .with_tool_name(tool_name.clone())
285 .with_execution_index(idx),
286 })
287 })
288 .collect()
289}