systemprompt_agent/services/
message.rs1use crate::services::shared::{AgentServiceError, Result};
8use serde_json::json;
9use uuid::Uuid;
10
11use crate::models::a2a::{Message, MessageRole, Part, TextPart};
12use crate::repository::context::message::PersistMessageWithTxParams;
13use crate::repository::task::TaskRepository;
14use systemprompt_database::{DatabaseProvider, DatabaseTransaction};
15use systemprompt_identifiers::{ContextId, MessageId, TaskId};
16use systemprompt_models::RequestContext;
17
18pub struct PersistMessageInTxParams<'a> {
19 pub tx: &'a mut dyn DatabaseTransaction,
20 pub message: &'a Message,
21 pub task_id: &'a TaskId,
22 pub context_id: &'a ContextId,
23 pub user_id: Option<&'a systemprompt_identifiers::UserId>,
24 pub session_id: &'a systemprompt_identifiers::SessionId,
25 pub trace_id: &'a systemprompt_identifiers::TraceId,
26}
27
28impl std::fmt::Debug for PersistMessageInTxParams<'_> {
29 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
30 f.debug_struct("PersistMessageInTxParams")
31 .field("message", &self.message)
32 .field("task_id", &self.task_id)
33 .field("context_id", &self.context_id)
34 .field("user_id", &self.user_id)
35 .field("session_id", &self.session_id)
36 .field("trace_id", &self.trace_id)
37 .finish_non_exhaustive()
38 }
39}
40
41#[derive(Debug)]
42pub struct PersistMessagesParams<'a> {
43 pub task_id: &'a TaskId,
44 pub context_id: &'a ContextId,
45 pub messages: Vec<Message>,
46 pub user_id: Option<&'a systemprompt_identifiers::UserId>,
47 pub session_id: &'a systemprompt_identifiers::SessionId,
48 pub trace_id: &'a systemprompt_identifiers::TraceId,
49}
50
51#[derive(Debug)]
52pub struct CreateToolExecutionMessageParams<'a> {
53 pub task_id: &'a TaskId,
54 pub context_id: &'a ContextId,
55 pub tool_name: &'a str,
56 pub tool_args: &'a serde_json::Value,
57 pub request_context: &'a RequestContext,
58}
59
60pub struct MessageService {
61 task_repo: TaskRepository,
62}
63
64impl std::fmt::Debug for MessageService {
65 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
66 f.debug_struct("MessageService").finish_non_exhaustive()
67 }
68}
69
70impl MessageService {
71 #[must_use]
72 pub const fn new(task_repo: TaskRepository) -> Self {
73 Self { task_repo }
74 }
75
76 pub async fn persist_message_in_tx(&self, params: PersistMessageInTxParams<'_>) -> Result<i32> {
77 let PersistMessageInTxParams {
78 tx,
79 message,
80 task_id,
81 context_id,
82 user_id,
83 session_id,
84 trace_id,
85 } = params;
86 let sequence_number = self
87 .task_repo
88 .get_next_sequence_number_in_tx(tx, task_id)
89 .await?;
90
91 self.task_repo
92 .persist_message_with_tx(PersistMessageWithTxParams {
93 tx,
94 message,
95 task_id,
96 context_id,
97 sequence_number,
98 user_id,
99 session_id,
100 trace_id,
101 })
102 .await
103 .map_err(|e| {
104 AgentServiceError::Internal(format!("Failed to persist message: {}", e))
105 })?;
106
107 tracing::info!(
108 message_id = %message.message_id,
109 task_id = %task_id,
110 sequence_number = sequence_number,
111 "Message persisted"
112 );
113
114 Ok(sequence_number)
115 }
116
117 pub async fn persist_messages(&self, params: PersistMessagesParams<'_>) -> Result<Vec<i32>> {
118 let PersistMessagesParams {
119 task_id,
120 context_id,
121 messages,
122 user_id,
123 session_id,
124 trace_id,
125 } = params;
126
127 if messages.is_empty() {
128 return Ok(Vec::new());
129 }
130
131 let mut tx = self
132 .task_repo
133 .db_pool()
134 .as_ref()
135 .begin_transaction()
136 .await?;
137 let mut sequence_numbers = Vec::new();
138
139 tracing::info!(
140 task_id = %task_id,
141 message_count = messages.len(),
142 "Persisting multiple messages"
143 );
144
145 for message in messages {
146 let seq = self
147 .persist_message_in_tx(PersistMessageInTxParams {
148 tx: &mut *tx,
149 message: &message,
150 task_id,
151 context_id,
152 user_id,
153 session_id,
154 trace_id,
155 })
156 .await?;
157 sequence_numbers.push(seq);
158 }
159
160 tx.commit().await?;
161
162 tracing::info!(
163 task_id = %task_id,
164 sequence_numbers = ?sequence_numbers,
165 "Messages persisted successfully"
166 );
167
168 Ok(sequence_numbers)
169 }
170
171 pub async fn create_tool_execution_message(
172 &self,
173 params: CreateToolExecutionMessageParams<'_>,
174 ) -> Result<(String, i32)> {
175 let CreateToolExecutionMessageParams {
176 task_id,
177 context_id,
178 tool_name,
179 tool_args,
180 request_context,
181 } = params;
182 let message_id = Uuid::new_v4().to_string();
183
184 let tool_args_display =
185 serde_json::to_string_pretty(tool_args).unwrap_or_else(|_| tool_args.to_string());
186
187 let timestamp = chrono::Utc::now().to_rfc3339();
188
189 let message = Message {
190 role: MessageRole::User,
191 message_id: MessageId::new(message_id.clone()),
192 task_id: Some(task_id.clone()),
193 context_id: context_id.clone(),
194 parts: vec![Part::Text(TextPart {
195 text: format!(
196 "Executed MCP tool: {} with arguments:\n{}\n\nExecution ID: {} at {}",
197 tool_name,
198 tool_args_display,
199 task_id.as_str(),
200 timestamp
201 ),
202 })],
203 metadata: Some(json!({
204 "source": "mcp_direct_call",
205 "tool_name": tool_name,
206 "is_synthetic": true,
207 "tool_args": tool_args,
208 "execution_timestamp": timestamp,
209 })),
210 extensions: None,
211 reference_task_ids: None,
212 };
213
214 let mut tx = self
215 .task_repo
216 .db_pool()
217 .as_ref()
218 .begin_transaction()
219 .await?;
220
221 let sequence_number = self
222 .persist_message_in_tx(PersistMessageInTxParams {
223 tx: &mut *tx,
224 message: &message,
225 task_id,
226 context_id,
227 user_id: Some(request_context.user_id()),
228 session_id: request_context.session_id(),
229 trace_id: request_context.trace_id(),
230 })
231 .await?;
232
233 tx.commit().await?;
234
235 tracing::info!(
236 message_id = %message_id,
237 task_id = %task_id,
238 tool_name = %tool_name,
239 sequence_number = sequence_number,
240 "Created synthetic tool execution message"
241 );
242
243 Ok((message_id, sequence_number))
244 }
245}