Skip to main content

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

1use 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}