systemprompt_agent/services/a2a_server/processing/
message_validation.rs

1use crate::services::shared::{AgentServiceError, Result};
2use systemprompt_database::DbPool;
3use systemprompt_identifiers::{ContextId, TaskId};
4use systemprompt_models::RequestContext;
5
6use crate::models::{AgentRuntimeInfo, Message};
7use crate::repository::context::ContextRepository;
8use crate::services::registry::AgentRegistry;
9
10#[derive(Clone, Debug)]
11pub struct ValidatedMessageRequest {
12    pub message: Message,
13    pub agent_name: String,
14    pub context_id: ContextId,
15    pub task_id: TaskId,
16    pub agent_runtime: AgentRuntimeInfo,
17    pub has_tools: bool,
18}
19
20#[derive(Debug)]
21pub struct MessageValidationService {
22    db_pool: DbPool,
23}
24
25impl MessageValidationService {
26    pub const fn new(db_pool: DbPool) -> Self {
27        Self { db_pool }
28    }
29
30    pub async fn validate_message_request(
31        &self,
32        message: &Message,
33        agent_name: &str,
34        context: &RequestContext,
35    ) -> Result<ValidatedMessageRequest> {
36        Self::validate_message_format(message)?;
37
38        let agent_runtime = self.load_agent_runtime(agent_name).await?;
39
40        self.validate_context_ownership(message, context).await?;
41
42        let task_id = Self::determine_task_id(message);
43
44        let has_tools = !agent_runtime.mcp_servers.include.is_empty();
45
46        let context_id = message.context_id.clone();
47
48        Ok(ValidatedMessageRequest {
49            message: message.clone(),
50            agent_name: agent_name.to_owned(),
51            context_id,
52            task_id,
53            agent_runtime,
54            has_tools,
55        })
56    }
57
58    async fn load_agent_runtime(&self, agent_name: &str) -> Result<AgentRuntimeInfo> {
59        let registry = AgentRegistry::new()?;
60        let agent_config = registry.get_agent(agent_name).await.map_err(|_e| {
61            AgentServiceError::Internal(format!("Agent not found: {}", agent_name))
62        })?;
63
64        Ok(agent_config.into())
65    }
66
67    async fn validate_context_ownership(
68        &self,
69        message: &Message,
70        context: &RequestContext,
71    ) -> Result<()> {
72        let context_repo = ContextRepository::new(&self.db_pool)?;
73
74        let context_id = &message.context_id;
75
76        context_repo
77            .get_context(context_id, context.user_id())
78            .await
79            .map_err(|e| {
80                AgentServiceError::Internal(format!(
81                    "Context validation failed - context_id: {}, user_id: {}, error: {}",
82                    context_id,
83                    context.user_id(),
84                    e
85                ))
86            })?;
87
88        Ok(())
89    }
90
91    #[doc(hidden)]
92    pub fn validate_message_format(message: &Message) -> Result<()> {
93        let has_text_part = message
94            .parts
95            .iter()
96            .any(|part| matches!(part, crate::models::Part::Text(_)));
97
98        if !has_text_part {
99            return Err(AgentServiceError::Internal(
100                "No text content found in message".to_owned(),
101            ));
102        }
103
104        Ok(())
105    }
106
107    #[doc(hidden)]
108    pub fn determine_task_id(message: &Message) -> TaskId {
109        message
110            .task_id
111            .clone()
112            .unwrap_or_else(|| TaskId::new(uuid::Uuid::new_v4().to_string()))
113    }
114}
systemprompt_agent/services/a2a_server/processing/message_validation.rs

systemprompt_agent/services/a2a_server/processing/
message_validation.rs
