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//!
4//! Copyright (c) systemprompt.io — Business Source License 1.1.
5//! See <https://systemprompt.io> for licensing details.
6
7use std::sync::Arc;
8
9use axum::response::sse::Event;
10use systemprompt_identifiers::{ContextId, MessageId, TaskId};
11use systemprompt_models::{
12    A2AEventBuilder, AgUiEventBuilder, CallToolResult, RequestContext, ToolCall,
13};
14use tokio::sync::mpsc::{Receiver, Sender};
15
16use crate::models::ExecutionStep;
17use crate::models::a2a::jsonrpc::NumberOrString;
18use crate::models::a2a::{Artifact, Message, TaskState};
19use crate::repository::task::TaskRepository;
20use crate::services::a2a_server::processing::message::{MessageProcessor, StreamEvent};
21
22use super::event_loop_lifecycle::{
23    EmitRunStartedParams, SendA2aStatusEventParams, emit_run_started, send_a2a_status_event,
24};
25use super::handlers::{
26    HandleCompleteParams, HandleErrorParams, TextStreamState, handle_complete, handle_error,
27};
28use super::webhook_client::WebhookContext;
29
30pub struct ProcessEventsParams {
31    pub tx: Sender<Event>,
32    pub chunk_rx: Receiver<StreamEvent>,
33    pub task_id: TaskId,
34    pub context_id: ContextId,
35    pub message_id: MessageId,
36    pub original_message: Message,
37    pub agent_name: String,
38    pub context: RequestContext,
39    pub task_repo: TaskRepository,
40    pub processor: Arc<MessageProcessor>,
41    pub request_id: NumberOrString,
42}
43
44impl std::fmt::Debug for ProcessEventsParams {
45    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
46        f.debug_struct("ProcessEventsParams")
47            .field("task_id", &self.task_id)
48            .field("context_id", &self.context_id)
49            .field("message_id", &self.message_id)
50            .field("agent_name", &self.agent_name)
51            .finish_non_exhaustive()
52    }
53}
54
55pub async fn process_events(params: ProcessEventsParams) {
56    let ProcessEventsParams {
57        tx,
58        mut chunk_rx,
59        task_id,
60        context_id,
61        message_id,
62        original_message,
63        agent_name,
64        context,
65        task_repo,
66        processor,
67        request_id,
68    } = params;
69
70    let webhook_context =
71        WebhookContext::new(context.user_id().clone(), context.auth_token().as_str());
72
73    emit_run_started(EmitRunStartedParams {
74        tx: &tx,
75        webhook_context: &webhook_context,
76        context_id: &context_id,
77        task_id: &task_id,
78        task_repo: &task_repo,
79        request_id: &request_id,
80    })
81    .await;
82
83    tracing::info!("Stream channel received, waiting for events...");
84
85    let mut text_state = TextStreamState::new().with_webhook_context(webhook_context.clone());
86
87    let ctx = EventLoopCtx {
88        tx: &tx,
89        webhook_context: &webhook_context,
90        task_id: &task_id,
91        context_id: &context_id,
92        message_id: &message_id,
93        original_message: &original_message,
94        agent_name: &agent_name,
95        context: &context,
96        task_repo: &task_repo,
97        processor: &processor,
98        request_id: &request_id,
99    };
100
101    while let Some(event) = chunk_rx.recv().await {
102        match event {
103            StreamEvent::Text(text) => {
104                text_state.handle_text(text, &message_id).await;
105            },
106            StreamEvent::ToolCallStarted(tool_call) => {
107                broadcast_tool_call_started(&webhook_context, &tool_call, &message_id).await;
108            },
109            StreamEvent::ToolResult { call_id, result } => {
110                broadcast_tool_result(&webhook_context, &call_id, &result).await;
111            },
112            StreamEvent::ExecutionStepUpdate { step } => {
113                broadcast_execution_step(&webhook_context, step, &context_id).await;
114            },
115            StreamEvent::Complete {
116                full_text,
117                artifacts,
118            } => {
119                text_state.finalize(&message_id).await;
120                finish_completed(&ctx, full_text, artifacts).await;
121                break;
122            },
123            StreamEvent::Error(error) => {
124                text_state.finalize(&message_id).await;
125                finish_failed(&ctx, error).await;
126                break;
127            },
128        }
129    }
130
131    drop(tx);
132
133    tracing::info!("Stream event loop ended");
134}
135
136struct EventLoopCtx<'a> {
137    tx: &'a Sender<Event>,
138    webhook_context: &'a WebhookContext,
139    task_id: &'a TaskId,
140    context_id: &'a ContextId,
141    message_id: &'a MessageId,
142    original_message: &'a Message,
143    agent_name: &'a str,
144    context: &'a RequestContext,
145    task_repo: &'a TaskRepository,
146    processor: &'a Arc<MessageProcessor>,
147    request_id: &'a NumberOrString,
148}
149
150async fn broadcast_tool_call_started(
151    webhook_context: &WebhookContext,
152    tool_call: &ToolCall,
153    message_id: &MessageId,
154) {
155    let tool_call_id = tool_call.ai_tool_call_id.as_str();
156    let start_event = AgUiEventBuilder::tool_call_start(
157        tool_call_id,
158        &tool_call.name,
159        Some(message_id.to_string()),
160    );
161    if let Err(e) = webhook_context.broadcast_agui(start_event).await {
162        tracing::error!(error = %e, "Failed to broadcast TOOL_CALL_START");
163    }
164
165    let args_json = serde_json::to_string(&tool_call.arguments).unwrap_or_else(|_| String::new());
166    let args_event = AgUiEventBuilder::tool_call_args(tool_call_id, &args_json);
167    if let Err(e) = webhook_context.broadcast_agui(args_event).await {
168        tracing::error!(error = %e, "Failed to broadcast TOOL_CALL_ARGS");
169    }
170
171    let end_event = AgUiEventBuilder::tool_call_end(tool_call_id);
172    if let Err(e) = webhook_context.broadcast_agui(end_event).await {
173        tracing::error!(error = %e, "Failed to broadcast TOOL_CALL_END");
174    }
175}
176
177async fn broadcast_tool_result(
178    webhook_context: &WebhookContext,
179    call_id: &str,
180    result: &CallToolResult,
181) {
182    let result_value = serde_json::to_value(result).unwrap_or_else(|_| serde_json::Value::Null);
183    let result_event =
184        AgUiEventBuilder::tool_call_result(uuid::Uuid::new_v4().to_string(), call_id, result_value);
185    if let Err(e) = webhook_context.broadcast_agui(result_event).await {
186        tracing::error!(error = %e, "Failed to broadcast TOOL_CALL_RESULT");
187    }
188}
189
190async fn broadcast_execution_step(
191    webhook_context: &WebhookContext,
192    step: ExecutionStep,
193    context_id: &ContextId,
194) {
195    let step_event = AgUiEventBuilder::execution_step(step, context_id.clone());
196    if let Err(e) = webhook_context.broadcast_agui(step_event).await {
197        tracing::error!(error = %e, "Failed to broadcast execution_step");
198    }
199}
200
201async fn finish_completed(ctx: &EventLoopCtx<'_>, full_text: String, artifacts: Vec<Artifact>) {
202    let complete_params = HandleCompleteParams {
203        tx: ctx.tx,
204        webhook_context: ctx.webhook_context,
205        full_text,
206        artifacts,
207        task_id: ctx.task_id,
208        context_id: ctx.context_id,
209        id: ctx.message_id.as_str(),
210        original_message: ctx.original_message,
211        agent_name: ctx.agent_name,
212        context: ctx.context,
213        auth_token: ctx.context.auth_token().as_str(),
214        task_repo: ctx.task_repo,
215        processor: ctx.processor,
216    };
217    handle_complete(complete_params).await;
218
219    send_a2a_status_event(&SendA2aStatusEventParams {
220        tx: ctx.tx,
221        task_id: ctx.task_id,
222        context_id: ctx.context_id,
223        state: "completed",
224        is_final: true,
225        request_id: ctx.request_id,
226    });
227
228    let a2a_event = A2AEventBuilder::task_status_update(
229        ctx.task_id.clone(),
230        ctx.context_id.clone(),
231        TaskState::Completed,
232        None,
233    );
234    if let Err(e) = ctx.webhook_context.broadcast_a2a(a2a_event).await {
235        tracing::error!(error = %e, "Failed to broadcast A2A completed");
236    }
237}
238
239async fn finish_failed(ctx: &EventLoopCtx<'_>, error: String) {
240    handle_error(HandleErrorParams {
241        tx: ctx.tx,
242        webhook_context: ctx.webhook_context,
243        error,
244        task_id: ctx.task_id,
245        context_id: ctx.context_id,
246        task_repo: ctx.task_repo,
247    })
248    .await;
249
250    send_a2a_status_event(&SendA2aStatusEventParams {
251        tx: ctx.tx,
252        task_id: ctx.task_id,
253        context_id: ctx.context_id,
254        state: "failed",
255        is_final: true,
256        request_id: ctx.request_id,
257    });
258
259    let a2a_event = A2AEventBuilder::task_status_update(
260        ctx.task_id.clone(),
261        ctx.context_id.clone(),
262        TaskState::Failed,
263        None,
264    );
265    if let Err(e) = ctx.webhook_context.broadcast_a2a(a2a_event).await {
266        tracing::error!(error = %e, "Failed to broadcast A2A failed");
267    }
268}
systemprompt_agent/services/a2a_server/streaming/event_loop.rs

systemprompt_agent/services/a2a_server/streaming/
event_loop.rs
