Skip to main content

adk_server/a2a/
processor.rs

1use crate::a2a::{
2    Artifact, TaskArtifactUpdateEvent, TaskState, TaskStatus, TaskStatusUpdateEvent,
3    metadata::{InvocationMeta, to_event_meta},
4    parts::adk_parts_to_a2a,
5};
6use adk_core::{Event, EventActions, Result};
7
8pub struct EventProcessor {
9    context_id: String,
10    task_id: String,
11    meta: InvocationMeta,
12    terminal_actions: EventActions,
13    response_id: Option<String>,
14    terminal_state: Option<TaskState>,
15    has_artifacts: bool,
16    text_deltas: adk_core::EventTextDeltas,
17}
18
19impl EventProcessor {
20    pub fn new(context_id: String, task_id: String, meta: InvocationMeta) -> Self {
21        Self {
22            context_id,
23            task_id,
24            meta,
25            terminal_actions: EventActions::default(),
26            response_id: None,
27            terminal_state: None,
28            has_artifacts: false,
29            text_deltas: adk_core::EventTextDeltas::default(),
30        }
31    }
32
33    pub fn process(&mut self, event: &Event) -> Result<Option<TaskArtifactUpdateEvent>> {
34        self.update_terminal_actions(event);
35        let event = self.text_deltas.push(event);
36
37        let event_meta = to_event_meta(&self.meta, &event);
38        let event_meta_map: serde_json::Map<String, serde_json::Value> =
39            event_meta.into_iter().collect();
40
41        // Get content
42        let content = match &event.llm_response.content {
43            Some(c) => c,
44            None => return Ok(None),
45        };
46
47        if content.parts.is_empty() {
48            return Ok(None);
49        }
50
51        // Convert parts
52        let parts = adk_parts_to_a2a(&content.parts, &[])?;
53
54        if parts.is_empty() {
55            return Ok(None);
56        }
57
58        self.has_artifacts = true;
59
60        let artifact_event = if let Some(response_id) = &self.response_id {
61            TaskArtifactUpdateEvent {
62                task_id: self.task_id.clone(),
63                context_id: Some(self.context_id.clone()),
64                artifact: Artifact {
65                    artifact_id: response_id.clone(),
66                    name: None,
67                    description: None,
68                    parts,
69                    metadata: Some(event_meta_map),
70                    extensions: None,
71                },
72                append: true,
73                last_chunk: false,
74            }
75        } else {
76            let artifact_id = uuid::Uuid::new_v4().to_string();
77            self.response_id = Some(artifact_id.clone());
78
79            TaskArtifactUpdateEvent {
80                task_id: self.task_id.clone(),
81                context_id: Some(self.context_id.clone()),
82                artifact: Artifact {
83                    artifact_id,
84                    name: None,
85                    description: None,
86                    parts,
87                    metadata: Some(event_meta_map),
88                    extensions: None,
89                },
90                append: true,
91                last_chunk: false,
92            }
93        };
94
95        Ok(Some(artifact_event))
96    }
97
98    pub fn make_terminal_events(&self) -> Vec<TaskStatusUpdateEvent> {
99        let mut events = vec![];
100
101        // Terminal status
102        let state = self.terminal_state.clone().unwrap_or(TaskState::Completed);
103
104        events.push(TaskStatusUpdateEvent {
105            task_id: self.task_id.clone(),
106            context_id: Some(self.context_id.clone()),
107            status: TaskStatus { state, message: None },
108            final_update: true,
109        });
110
111        events
112    }
113
114    fn update_terminal_actions(&mut self, event: &Event) {
115        self.terminal_actions.escalate = self.terminal_actions.escalate || event.actions.escalate;
116        if let Some(agent) = &event.actions.transfer_to_agent {
117            self.terminal_actions.transfer_to_agent = Some(agent.clone());
118        }
119    }
120}