1use async_trait::async_trait;
2use cookie_agent_engine::{
3 AgentMessageInvocation, Engine, MESSAGE_INVALID_ARGUMENTS, MESSAGE_INVALID_BODY,
4 PreparedExecutor, PreparedTool, SessionToolContext, SubagentScope, ToolCall, ToolCompletion,
5 ToolError, ToolExecutionContext, ToolPreparationContext, ToolProvider, ToolSpec,
6};
7use cookie_agent_protocol::{
8 PermissionAction, PersistedToolResult as ToolResult, ProducerDeliveryMode, SessionId,
9};
10use serde::{Deserialize, Serialize};
11
12use crate::{fs_cap, prepared_operation, prepared_resource, safe_title};
13
14fn default_mode() -> ProducerDeliveryMode {
15 ProducerDeliveryMode::Steer
16}
17
18#[derive(Debug, Deserialize, Serialize)]
19#[serde(deny_unknown_fields)]
20struct MessageArgs {
21 to: String,
23 body: String,
24 #[serde(default = "default_mode")]
25 mode: ProducerDeliveryMode,
26}
27
28pub struct MessageToolProvider {
29 engine: Engine,
30}
31
32impl MessageToolProvider {
33 #[must_use]
34 pub fn new(engine: Engine) -> Self {
35 Self { engine }
36 }
37
38 fn spec() -> ToolSpec {
43 ToolSpec {
44 output: Default::default(),
45 concurrency: cookie_agent_engine::ToolConcurrency::Parallel,
46 result_truncation: cookie_agent_engine::ToolResultTruncationPolicy::Bounded,
47 name: "send_message".into(),
48 permission_name: "message".into(),
49 description: "Send a message to another agent in your session tree. Recipient accepts a subagent handle or a full UUID.".into(),
50 parameters: serde_json::json!({
51 "type":"object","additionalProperties":false,
52 "properties":{
53 "to":{"type":"string","description":"Handle or full UUID of the recipient agent."},
54 "body":{"type":"string","minLength":1},
55 "mode":{"type":"string","enum":["steer","queue"],"default":"steer"}
56 },"required":["to","body"]
57 }),
58 }
59 }
60
61 fn parse(args: &serde_json::Value) -> Result<MessageArgs, ToolError> {
64 let value: MessageArgs = serde_json::from_value(args.clone())
65 .map_err(|_| ToolError::execution(MESSAGE_INVALID_ARGUMENTS))?;
66 if value.body.trim().is_empty() {
67 return Err(ToolError::execution(MESSAGE_INVALID_BODY));
68 }
69 Ok(value)
70 }
71
72 fn resolve_recipient(
77 engine: &Engine,
78 sender: SessionId,
79 reference: &str,
80 ) -> Result<SessionId, ToolError> {
81 if let Ok(id) = reference.parse::<SessionId>() {
82 return Ok(id);
83 }
84 engine
85 .resolve_subagent_target(sender, reference, SubagentScope::TreePeers)
86 .map_err(|error| ToolError::execution(error.to_string()))
87 }
88}
89
90#[async_trait]
91impl ToolProvider for MessageToolProvider {
92 fn provider_id(&self) -> &'static str {
93 "builtin.message"
94 }
95
96 fn tools_for_session(&self, _ctx: &SessionToolContext) -> Result<Vec<ToolSpec>, ToolError> {
97 Ok(vec![Self::spec()])
98 }
99
100 fn get_permission_name(_name: &str) -> Result<&'static str, ToolError> {
101 Ok("message")
102 }
103
104 fn get_permission_resource(
105 &self,
106 _name: &str,
107 arguments: &serde_json::Value,
108 ) -> Result<(&'static str, Option<String>), ToolError> {
109 let _ = Self::parse(arguments)?;
110 Ok(("message", Some("*".into())))
111 }
112
113 fn get_display_argument(
114 &self,
115 _name: &str,
116 arguments: &serde_json::Value,
117 ) -> Result<String, ToolError> {
118 Ok(Self::parse(arguments)?.to.to_string())
119 }
120
121 async fn prepare(
122 &self,
123 ctx: ToolPreparationContext,
124 call: ToolCall,
125 ) -> Result<PreparedTool, ToolError> {
126 if call.name != "send_message" {
127 return Err(ToolError::execution(
128 "message provider received another tool",
129 ));
130 }
131 let args = Self::parse(&call.arguments)?;
132 let recipient = Self::resolve_recipient(&self.engine, ctx.session, &args.to)?;
133 let relationship = self
134 .engine
135 .message_relationship(ctx.session, recipient)
136 .map_err(|error| ToolError::execution(error.to_string()))?;
137 let cwd = fs_cap::cwd_context_bytes(&ctx.cwd)?;
138 let operation = prepared_operation(
139 "send_message",
140 &args,
141 vec![(PermissionAction::Message, relationship)],
142 vec![prepared_resource(
143 PermissionAction::Message,
144 relationship,
145 relationship.as_bytes(),
146 relationship.as_bytes(),
147 cookie_agent_protocol::PreparedBindingLifetime::RestartStable,
148 cookie_agent_protocol::ApprovalResourceSource::PrimaryOperation,
149 )?],
150 &cwd,
151 )?;
152 let normalized =
153 serde_json::to_value(&args).map_err(|error| ToolError::execution(error.to_string()))?;
154 PreparedTool::new(
155 operation,
156 normalized,
157 None,
158 Box::new(MessageExecutor {
159 engine: self.engine.clone(),
160 args,
161 call_id: call.id,
162 }),
163 )?
164 .with_policy_labels(vec![relationship.into()])
165 }
166}
167
168struct MessageExecutor {
169 engine: Engine,
170 args: MessageArgs,
171 call_id: cookie_agent_protocol::ToolCallId,
172}
173
174#[async_trait]
175impl PreparedExecutor for MessageExecutor {
176 async fn revalidate(&self) -> Result<(), ToolError> {
177 Ok(())
178 }
179 async fn execute(
180 self: Box<Self>,
181 context: ToolExecutionContext,
182 ) -> Result<ToolCompletion, ToolError> {
183 let recipient =
184 MessageToolProvider::resolve_recipient(&self.engine, context.session, &self.args.to)?;
185 let handle = self
186 .engine
187 .send_agent_message(AgentMessageInvocation {
188 sender_session_id: context.session,
189 sender_run_id: context.run,
190 sender_tool_call_id: self.call_id,
191 recipient_session_id: recipient,
192 body: self.args.body,
193 mode: self.args.mode,
194 })
195 .await
196 .map_err(|error| ToolError::execution(error.to_string()))?;
197 let metadata = serde_json::json!({"message_id": handle.message_id, "mode": handle.mode, "recipient_state": handle.recipient_state.as_str()});
198 Ok(ToolCompletion::single(ToolResult {
199 display: None,
200 retained_output: None,
201 title: safe_title("Message sent"),
202 output: format!("Message accepted for {}.", handle.recipient_state.as_str()),
203 metadata,
204 truncation: None,
205 attachments: Vec::new(),
206 additional_messages: Vec::new(),
207 }))
208 }
209}
210
211#[cfg(test)]
212mod tests {
213 use cookie_agent_protocol::SessionId;
214
215 use super::{MessageArgs, MessageToolProvider};
216
217 fn arguments(body: &str) -> serde_json::Value {
218 serde_json::json!({"to": SessionId::new_v7().to_string(), "body": body})
219 }
220
221 fn rejection(arguments: &serde_json::Value) -> String {
223 let rendered = MessageToolProvider::parse(arguments)
224 .expect_err("arguments must be rejected")
225 .to_string();
226 rendered
227 .strip_prefix("tool failed: ")
228 .unwrap_or(&rendered)
229 .to_owned()
230 }
231
232 #[test]
233 fn surface_addresses_the_recipient_with_to() {
234 let spec = MessageToolProvider::spec().parameters;
235 assert_eq!(spec["properties"]["to"]["type"], "string");
236 assert_eq!(spec["required"], serde_json::json!(["to", "body"]));
237 assert_eq!(spec["additionalProperties"], false);
238 assert!(spec["properties"].get("recipient_session_id").is_none());
239 }
240
241 #[test]
242 fn parse_accepts_the_documented_shape_and_defaults_the_mode() {
243 let recipient = SessionId::new_v7();
244 let parsed = MessageToolProvider::parse(&serde_json::json!({
245 "to": recipient.to_string(),
246 "body":"chain start"
247 }))
248 .expect("valid arguments");
249 assert_eq!(parsed.to, recipient.to_string());
250 assert_eq!(parsed.body, "chain start");
251 assert_eq!(
252 parsed.mode,
253 cookie_agent_protocol::ProducerDeliveryMode::Steer
254 );
255 let queued = MessageToolProvider::parse(&serde_json::json!({
256 "to": recipient.to_string(),
257 "body":"chain start",
258 "mode":"queue"
259 }))
260 .expect("queued arguments");
261 assert_eq!(
262 queued.mode,
263 cookie_agent_protocol::ProducerDeliveryMode::Queue
264 );
265 }
266
267 #[test]
268 fn normalized_arguments_record_the_to_field_only() {
269 let recipient = SessionId::new_v7();
270 let parsed = MessageToolProvider::parse(&serde_json::json!({
271 "to": recipient.to_string(),
272 "body":"chain start"
273 }))
274 .expect("valid arguments");
275 let normalized = serde_json::to_value(&parsed).expect("normalized arguments");
276 assert_eq!(normalized["to"], recipient.to_string());
277 assert!(normalized.get("recipient_session_id").is_none());
278 }
279
280 #[test]
281 fn every_malformed_shape_reports_one_stable_code() {
282 let recipient = SessionId::new_v7();
283 for arguments in [
284 &serde_json::json!({"recipient_session_id": recipient.to_string(), "body":"mail"}),
286 &serde_json::json!({"body":"missing recipient"}),
287 &serde_json::json!({"to":recipient.to_string(),"body":"mail","extra":true}),
290 &serde_json::json!({"to":recipient.to_string(),"body":42}),
291 &serde_json::json!({"to":recipient.to_string(),"body":"mail","mode":"later"}),
292 &serde_json::json!([]),
293 ] {
294 assert_eq!(
295 rejection(arguments),
296 cookie_agent_engine::MESSAGE_INVALID_ARGUMENTS,
297 "unstable parse error for {arguments}"
298 );
299 }
300 }
301
302 #[test]
303 fn blank_bodies_keep_their_own_stable_code() {
304 for body in ["", " ", "\n\t"] {
305 assert_eq!(
306 rejection(&arguments(body)),
307 cookie_agent_engine::MESSAGE_INVALID_BODY
308 );
309 }
310 }
311
312 #[test]
313 fn argument_shape_stays_private_to_the_tool() {
314 let normalized = serde_json::to_value(MessageArgs {
317 to: SessionId::new_v7().to_string(),
318 body: "mail".into(),
319 mode: cookie_agent_protocol::ProducerDeliveryMode::Queue,
320 })
321 .expect("serializable");
322 let mut keys = normalized
323 .as_object()
324 .expect("object")
325 .keys()
326 .cloned()
327 .collect::<Vec<_>>();
328 keys.sort();
329 assert_eq!(keys, vec!["body", "mode", "to"]);
330 }
331}