Skip to main content

cookie_agent_tools/
message.rs

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    /// Handle or full UUID of the recipient.
22    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    /// The frozen `send_message` surface. The recipient argument is `to`; there
39    /// is deliberately no compatibility alias for the earlier
40    /// `recipient_session_id` name, which would otherwise be frozen into every
41    /// normalized-arguments record this call produces.
42    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    /// Parses one call, mapping every malformed shape onto a stable code so the
62    /// text a sending model sees never depends on serde's message formatting.
63    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    /// Resolves the recipient: a full UUID passes through unchanged (so the
73    /// engine's `not_tree_peer` / `self_send` guards keep their stable codes);
74    /// anything else is a handle resolved within the sender's tree, with the
75    /// self-repairing candidate-list error on failure.
76    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    /// The execution-message wrapper keeps the stable code verbatim.
222    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            // The pre-rename argument name is unknown, not an alias.
285            &serde_json::json!({"recipient_session_id": recipient.to_string(), "body":"mail"}),
286            &serde_json::json!({"body":"missing recipient"}),
287            // A non-UUID string is a loose-typed reference the resolver rejects
288            // later, not a parse error.
289            &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        // The executor forwards `to` as the engine's recipient field; the model
315        // never sees that name.
316        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}