Skip to main content

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

1use super::helpers::{content_to_json, extract_text_from_content};
2use super::TaskBuilder;
3use crate::models::a2a::{
4    Artifact, DataPart, Message, 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::{agent_names, ArtifactMetadata, TaskMetadata};
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        kind: "task".to_string(),
55        status: TaskStatus {
56            state: TaskState::Submitted,
57            message: None,
58            timestamp: Some(chrono::Utc::now()),
59        },
60        history: Some(vec![user_message]),
61        artifacts: None,
62        metadata: Some(TaskMetadata::new_agent_message(agent_name.to_string())),
63    }
64}
65
66pub fn build_multiturn_task(
67    context_id: ContextId,
68    task_id: TaskId,
69    user_message: Message,
70    tool_calls: Vec<ToolCall>,
71    tool_results: Vec<CallToolResult>,
72    final_response: String,
73    total_iterations: usize,
74) -> Task {
75    let ctx_id = context_id;
76
77    let history = build_history(
78        &ctx_id,
79        &task_id,
80        user_message,
81        &tool_calls,
82        &tool_results,
83        &final_response,
84    );
85
86    let artifacts = build_artifacts(&ctx_id, &task_id, &tool_calls, &tool_results);
87
88    Task {
89        id: task_id.clone(),
90        context_id: ctx_id.clone(),
91        kind: "task".to_string(),
92        status: TaskStatus {
93            state: TaskState::Completed,
94            message: Some(Message {
95                role: "agent".to_string(),
96                parts: vec![Part::Text(TextPart {
97                    text: final_response,
98                })],
99                id: MessageId::generate(),
100                task_id: Some(task_id),
101                context_id: ctx_id,
102                kind: "message".to_string(),
103                metadata: None,
104                extensions: None,
105                reference_task_ids: None,
106            }),
107            timestamp: Some(chrono::Utc::now()),
108        },
109        history: Some(history),
110        artifacts: if artifacts.is_empty() {
111            None
112        } else {
113            Some(artifacts)
114        },
115        metadata: Some(
116            TaskMetadata::new_agent_message(agent_names::SYSTEM.to_string())
117                .with_extension("total_iterations".to_string(), json!(total_iterations))
118                .with_extension("total_tools_called".to_string(), json!(tool_calls.len())),
119        ),
120    }
121}
122
123fn build_history(
124    ctx_id: &ContextId,
125    task_id: &TaskId,
126    user_message: Message,
127    tool_calls: &[ToolCall],
128    tool_results: &[CallToolResult],
129    final_response: &str,
130) -> Vec<Message> {
131    let mut history = Vec::new();
132    history.push(user_message);
133
134    let mut iteration = 1;
135    let mut call_idx = 0;
136
137    while call_idx < tool_calls.len() {
138        let iteration_calls: Vec<_> = tool_calls
139            .iter()
140            .skip(call_idx)
141            .take_while(|_| call_idx < tool_calls.len())
142            .cloned()
143            .collect();
144
145        if iteration_calls.is_empty() {
146            break;
147        }
148
149        history.push(Message {
150            role: "agent".to_string(),
151            parts: vec![Part::Text(TextPart {
152                text: format!("Executing {} tool(s)...", iteration_calls.len()),
153            })],
154            id: MessageId::generate(),
155            task_id: Some(task_id.clone()),
156            context_id: ctx_id.clone(),
157            kind: "message".to_string(),
158            metadata: Some(json!({
159                "iteration": iteration,
160                "tool_calls": iteration_calls.iter().map(|tc| {
161                    json!({"id": tc.ai_tool_call_id.as_ref(), "name": tc.name})
162                }).collect::<Vec<_>>()
163            })),
164            extensions: None,
165            reference_task_ids: None,
166        });
167
168        let results_text = iteration_calls
169            .iter()
170            .enumerate()
171            .filter_map(|(idx, call)| {
172                let result_idx = call_idx + idx;
173                tool_results.get(result_idx).map(|r| {
174                    let content_text = extract_text_from_content(&r.content);
175                    format!("Tool '{}' result: {}", call.name, content_text)
176                })
177            })
178            .collect::<Vec<_>>()
179            .join("\n");
180
181        history.push(Message {
182            role: "user".to_string(),
183            parts: vec![Part::Text(TextPart { text: results_text })],
184            id: MessageId::generate(),
185            task_id: Some(task_id.clone()),
186            context_id: ctx_id.clone(),
187            kind: "message".to_string(),
188            metadata: Some(json!({
189                "iteration": iteration,
190                "tool_results": true
191            })),
192            extensions: None,
193            reference_task_ids: None,
194        });
195
196        call_idx += iteration_calls.len();
197        iteration += 1;
198    }
199
200    history.push(Message {
201        role: "agent".to_string(),
202        parts: vec![Part::Text(TextPart {
203            text: final_response.to_string(),
204        })],
205        id: MessageId::generate(),
206        task_id: Some(task_id.clone()),
207        context_id: ctx_id.clone(),
208        kind: "message".to_string(),
209        metadata: Some(json!({
210            "iteration": iteration,
211            "final_synthesis": true
212        })),
213        extensions: None,
214        reference_task_ids: None,
215    });
216
217    history
218}
219
220fn build_artifacts(
221    ctx_id: &ContextId,
222    task_id: &TaskId,
223    tool_calls: &[ToolCall],
224    tool_results: &[CallToolResult],
225) -> Vec<Artifact> {
226    tool_results
227        .iter()
228        .enumerate()
229        .filter_map(|(idx, result)| {
230            let tool_call = tool_calls.get(idx)?;
231            let tool_name = &tool_call.name;
232            let call_id = tool_call.ai_tool_call_id.as_ref();
233            let is_error = result.is_error?;
234
235            let structured_content = result.structured_content.as_ref()?;
236            let parsed = parse_tool_response(structured_content)
237                .map_err(|e| {
238                    tracing::debug!(tool_name = %tool_name, error = %e, "Failed to parse tool response, skipping artifact");
239                    e
240                })
241                .ok()?;
242
243            let mut data_map = serde_json::Map::new();
244            data_map.insert("call_id".to_string(), json!(call_id));
245            data_map.insert("tool_name".to_string(), json!(tool_name));
246            data_map.insert("output".to_string(), content_to_json(&result.content));
247            data_map.insert(
248                "status".to_string(),
249                json!(if is_error { "error" } else { "success" }),
250            );
251
252            Some(Artifact {
253                id: parsed.artifact_id,
254                name: Some(format!("tool_execution_{}", idx + 1)),
255                description: Some(format!("Result from tool: {tool_name}")),
256                parts: vec![Part::Data(DataPart { data: data_map })],
257                extensions: vec![],
258                metadata: ArtifactMetadata::new(
259                    "tool_execution".to_string(),
260                    ctx_id.clone(),
261                    task_id.clone(),
262                )
263                .with_mcp_execution_id(call_id.to_string())
264                .with_tool_name(tool_name.to_string())
265                .with_execution_index(idx),
266            })
267        })
268        .collect()
269}