Skip to main content

systemprompt_agent/services/a2a_server/processing/message/
mod.rs

1//! Message processing for the A2A server.
2//!
3//! [`MessageProcessor`] owns the repositories and services needed to handle an
4//! inbound message and persist the resulting task. [`StreamProcessor`] drives
5//! the streaming execution pipeline, reporting progress as [`StreamEvent`]s
6//! over an mpsc channel. Both the streaming and non-streaming entry points live
7//! in the submodules.
8//!
9//! Copyright (c) systemprompt.io — Business Source License 1.1.
10//! See <https://systemprompt.io> for licensing details.
11
12mod message_handler;
13mod persistence;
14mod stream_processor;
15
16pub use stream_processor::StreamProcessor;
17
18use crate::services::shared::{AgentServiceError, Result};
19use std::sync::Arc;
20use tokio::sync::mpsc;
21
22use crate::models::AgentRuntimeInfo;
23use crate::models::a2a::{Artifact, Message, Task};
24use systemprompt_models::{AiProvider, CallToolResult, ToolCall};
25
26#[derive(Debug)]
27pub enum StreamEvent {
28    Text(String),
29    ToolCallStarted(ToolCall),
30    ToolResult {
31        call_id: String,
32        result: CallToolResult,
33    },
34    ExecutionStepUpdate {
35        step: crate::models::ExecutionStep,
36    },
37    Complete {
38        full_text: String,
39        artifacts: Vec<Artifact>,
40    },
41    Error(String),
42}
43use crate::repository::A2ARepositories;
44use crate::repository::execution::ExecutionStepRepository;
45use crate::services::{ContextService, SkillService};
46use systemprompt_identifiers::TaskId;
47use systemprompt_models::RequestContext;
48
49#[derive(Debug)]
50pub struct PersistCompletedTaskOnProcessorParams<'a> {
51    pub task: &'a Task,
52    pub user_message: &'a Message,
53    pub agent_message: &'a Message,
54    pub context: &'a RequestContext,
55    pub agent_name: &'a str,
56    pub artifacts_already_published: bool,
57}
58
59#[derive(Debug)]
60pub struct ProcessMessageStreamParams<'a> {
61    pub a2a_message: &'a Message,
62    pub agent_runtime: &'a AgentRuntimeInfo,
63    pub agent_name: &'a str,
64    pub context: &'a RequestContext,
65    pub task_id: TaskId,
66}
67
68pub struct MessageProcessor {
69    repositories: Arc<A2ARepositories>,
70    ai_service: Arc<dyn AiProvider>,
71    context_service: ContextService,
72    skill_service: Arc<SkillService>,
73    execution_step_repo: Arc<ExecutionStepRepository>,
74}
75
76impl std::fmt::Debug for MessageProcessor {
77    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
78        f.debug_struct("MessageProcessor")
79            .field("ai_service", &"<Arc<dyn AiProvider>>")
80            .finish()
81    }
82}
83
84impl MessageProcessor {
85    pub fn new(
86        repositories: Arc<A2ARepositories>,
87        ai_service: Arc<dyn AiProvider>,
88    ) -> Result<Self> {
89        let context_service = ContextService::new(repositories.tasks.clone());
90        let execution_step_repo = Arc::new(repositories.execution_steps.clone());
91        let skill_service = Arc::new(
92            SkillService::new()?.with_execution_step_repo(Arc::clone(&execution_step_repo)),
93        );
94
95        Ok(Self {
96            repositories,
97            ai_service,
98            context_service,
99            skill_service,
100            execution_step_repo,
101        })
102    }
103
104    pub async fn load_agent_runtime(&self, agent_name: &str) -> Result<AgentRuntimeInfo> {
105        use crate::services::registry::AgentRegistry;
106
107        let registry = AgentRegistry::new()?;
108        let agent_config = registry
109            .get_agent(agent_name)
110            .await
111            .map_err(|_e| AgentServiceError::Internal("Agent not found".to_owned()))?;
112
113        Ok(agent_config.into())
114    }
115
116    pub async fn persist_completed_task(
117        &self,
118        params: PersistCompletedTaskOnProcessorParams<'_>,
119    ) -> Result<Task> {
120        persistence::persist_completed_task(persistence::PersistCompletedTaskParams {
121            task: params.task,
122            user_message: params.user_message,
123            agent_message: params.agent_message,
124            context: params.context,
125            repositories: &self.repositories,
126            artifacts_already_published: params.artifacts_already_published,
127        })
128        .await
129    }
130
131    pub async fn process_message_stream(
132        &self,
133        params: ProcessMessageStreamParams<'_>,
134    ) -> Result<mpsc::Receiver<StreamEvent>> {
135        let stream_processor = StreamProcessor {
136            ai_service: Arc::clone(&self.ai_service),
137            context_service: self.context_service.clone(),
138            skill_service: Arc::clone(&self.skill_service),
139            execution_step_repo: Arc::clone(&self.execution_step_repo),
140        };
141
142        stream_processor.process_message_stream(params).await
143    }
144}