systemprompt_agent/services/a2a_server/streaming/
event_loop.rs

1//! Streaming event loop — fans `StreamEvent`s from the AI model out to A2A
2//! status frames, AG-UI webhooks, the SSE channel, and the task repository.
3
4use std::sync::Arc;
5
6use axum::response::sse::Event;
7use systemprompt_identifiers::{ContextId, MessageId, TaskId};
8use systemprompt_models::{A2AEventBuilder, AgUiEventBuilder, RequestContext};
9use tokio::sync::mpsc::{Receiver, Sender};
10
11use crate::models::a2a::jsonrpc::NumberOrString;
12use crate::models::a2a::{Message, TaskState};
13use crate::repository::task::TaskRepository;
14use crate::services::a2a_server::processing::message::{MessageProcessor, StreamEvent};
15
16use super::event_loop_lifecycle::{
17    EmitRunStartedParams, SendA2aStatusEventParams, emit_run_started, send_a2a_status_event,
18};
19use super::handlers::{
20    HandleCompleteParams, HandleErrorParams, TextStreamState, handle_complete, handle_error,
21};
22use super::webhook_client::WebhookContext;
23
24pub struct ProcessEventsParams {
25    pub tx: Sender<Event>,
26    pub chunk_rx: Receiver<StreamEvent>,
27    pub task_id: TaskId,
28    pub context_id: ContextId,
29    pub message_id: MessageId,
30    pub original_message: Message,
31    pub agent_name: String,
32    pub context: RequestContext,
33    pub task_repo: TaskRepository,
34    pub processor: Arc<MessageProcessor>,
35    pub request_id: NumberOrString,
36}
37
38impl std::fmt::Debug for ProcessEventsParams {
39    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40        f.debug_struct("ProcessEventsParams")
41            .field("task_id", &self.task_id)
42            .field("context_id", &self.context_id)
43            .field("message_id", &self.message_id)
44            .field("agent_name", &self.agent_name)
45            .finish_non_exhaustive()
46    }
47}
48
49pub(super) async fn process_events(params: ProcessEventsParams) {
50    let ProcessEventsParams {
51        tx,
52        mut chunk_rx,
53        task_id,
54        context_id,
55        message_id,
56        original_message,
57        agent_name,
58        context,
59        task_repo,
60        processor,
61        request_id,
62    } = params;
63
64    let webhook_context =
65        WebhookContext::new(context.user_id().clone(), context.auth_token().as_str());
66
67    emit_run_started(EmitRunStartedParams {
68        tx: &tx,
69        webhook_context: &webhook_context,
70        context_id: &context_id,
71        task_id: &task_id,
72        task_repo: &task_repo,
73        request_id: &request_id,
74    })
75    .await;
76
77    tracing::info!("Stream channel received, waiting for events...");
78
79    let mut text_state = TextStreamState::new().with_webhook_context(webhook_context.clone());
80
81    while let Some(event) = chunk_rx.recv().await {
82        match event {
83            StreamEvent::Text(text) => {
84                text_state.handle_text(text, &message_id).await;
85            },
86            StreamEvent::ToolCallStarted(tool_call) => {
87                let tool_call_id = tool_call.ai_tool_call_id.as_str();
88                let start_event = AgUiEventBuilder::tool_call_start(
89                    tool_call_id,
90                    &tool_call.name,
91                    Some(message_id.to_string()),
92                );
93                if let Err(e) = webhook_context.broadcast_agui(start_event).await {
94                    tracing::error!(error = %e, "Failed to broadcast TOOL_CALL_START");
95                }
96
97                let args_json =
98                    serde_json::to_string(&tool_call.arguments).unwrap_or_else(|_| String::new());
99                let args_event = AgUiEventBuilder::tool_call_args(tool_call_id, &args_json);
100                if let Err(e) = webhook_context.broadcast_agui(args_event).await {
101                    tracing::error!(error = %e, "Failed to broadcast TOOL_CALL_ARGS");
102                }
103
104                let end_event = AgUiEventBuilder::tool_call_end(tool_call_id);
105                if let Err(e) = webhook_context.broadcast_agui(end_event).await {
106                    tracing::error!(error = %e, "Failed to broadcast TOOL_CALL_END");
107                }
108            },
109            StreamEvent::ToolResult { call_id, result } => {
110                let result_value =
111                    serde_json::to_value(&result).unwrap_or_else(|_| serde_json::Value::Null);
112                let result_event = AgUiEventBuilder::tool_call_result(
113                    uuid::Uuid::new_v4().to_string(),
114                    &call_id,
115                    result_value,
116                );
117                if let Err(e) = webhook_context.broadcast_agui(result_event).await {
118                    tracing::error!(error = %e, "Failed to broadcast TOOL_CALL_RESULT");
119                }
120            },
121            StreamEvent::ExecutionStepUpdate { step } => {
122                let step_event = AgUiEventBuilder::execution_step(step.clone(), context_id.clone());
123                if let Err(e) = webhook_context.broadcast_agui(step_event).await {
124                    tracing::error!(error = %e, "Failed to broadcast execution_step");
125                }
126            },
127            StreamEvent::Complete {
128                full_text,
129                artifacts,
130            } => {
131                text_state.finalize(&message_id).await;
132
133                let complete_params = HandleCompleteParams {
134                    tx: &tx,
135                    webhook_context: &webhook_context,
136                    full_text,
137                    artifacts,
138                    task_id: &task_id,
139                    context_id: &context_id,
140                    id: message_id.as_str(),
141                    original_message: &original_message,
142                    agent_name: &agent_name,
143                    context: &context,
144                    auth_token: context.auth_token().as_str(),
145                    task_repo: &task_repo,
146                    processor: &processor,
147                };
148                handle_complete(complete_params).await;
149
150                send_a2a_status_event(&SendA2aStatusEventParams {
151                    tx: &tx,
152                    task_id: &task_id,
153                    context_id: &context_id,
154                    state: "completed",
155                    is_final: true,
156                    request_id: &request_id,
157                });
158
159                let a2a_event = A2AEventBuilder::task_status_update(
160                    task_id.clone(),
161                    context_id.clone(),
162                    TaskState::Completed,
163                    None,
164                );
165                if let Err(e) = webhook_context.broadcast_a2a(a2a_event).await {
166                    tracing::error!(error = %e, "Failed to broadcast A2A completed");
167                }
168
169                break;
170            },
171            StreamEvent::Error(error) => {
172                text_state.finalize(&message_id).await;
173                handle_error(HandleErrorParams {
174                    tx: &tx,
175                    webhook_context: &webhook_context,
176                    error,
177                    task_id: &task_id,
178                    context_id: &context_id,
179                    task_repo: &task_repo,
180                })
181                .await;
182
183                send_a2a_status_event(&SendA2aStatusEventParams {
184                    tx: &tx,
185                    task_id: &task_id,
186                    context_id: &context_id,
187                    state: "failed",
188                    is_final: true,
189                    request_id: &request_id,
190                });
191
192                let a2a_event = A2AEventBuilder::task_status_update(
193                    task_id.clone(),
194                    context_id.clone(),
195                    TaskState::Failed,
196                    None,
197                );
198                if let Err(e) = webhook_context.broadcast_a2a(a2a_event).await {
199                    tracing::error!(error = %e, "Failed to broadcast A2A failed");
200                }
201
202                break;
203            },
204        }
205    }
206
207    drop(tx);
208
209    tracing::info!("Stream event loop ended");
210}
systemprompt_agent/services/a2a_server/streaming/event_loop.rs

systemprompt_agent/services/a2a_server/streaming/
event_loop.rs
