Skip to main content

atman_runtime/
message.rs

1use std::path::PathBuf;
2
3use serde::{Deserialize, Serialize};
4
5use crate::event::TurnId;
6
7pub const TOOL_CALL_INTENT_FIELD: &str = "_atman_intent";
8pub const TOOL_CALL_INTENT_MAX_CHARS: usize = 120;
9
10#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
11#[serde(transparent)]
12pub struct ToolCallIntent(String);
13
14impl ToolCallIntent {
15    pub fn new(value: impl AsRef<str>) -> Option<Self> {
16        let normalized = value
17            .as_ref()
18            .split_whitespace()
19            .collect::<Vec<_>>()
20            .join(" ")
21            .chars()
22            .take(TOOL_CALL_INTENT_MAX_CHARS)
23            .collect::<String>();
24        (!normalized.is_empty()).then_some(Self(normalized))
25    }
26
27    pub fn as_str(&self) -> &str {
28        &self.0
29    }
30}
31
32impl<'de> Deserialize<'de> for ToolCallIntent {
33    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
34        let value = String::deserialize(deserializer)?;
35        Self::new(value).ok_or_else(|| serde::de::Error::custom("tool call intent is empty"))
36    }
37}
38
39pub fn decode_tool_call_input(
40    wire_input: serde_json::Value,
41    tool_name: &str,
42    tools: &[crate::tool::ToolSpec],
43) -> (serde_json::Value, Option<ToolCallIntent>) {
44    if !crate::tool::tool_spec_supports_call_intent(tool_name, tools) {
45        return (wire_input, None);
46    }
47    let serde_json::Value::Object(mut input) = wire_input else {
48        return (wire_input, None);
49    };
50    let intent = input
51        .remove(TOOL_CALL_INTENT_FIELD)
52        .and_then(|value| value.as_str().and_then(ToolCallIntent::new));
53    (serde_json::Value::Object(input), intent)
54}
55
56pub fn encode_tool_call_input(
57    clean_input: &serde_json::Value,
58    intent: Option<&ToolCallIntent>,
59    tool_name: &str,
60    tools: &[crate::tool::ToolSpec],
61) -> serde_json::Value {
62    let Some(intent) = intent else {
63        return clean_input.clone();
64    };
65    if crate::tool::tool_spec_blocks_call_intent(tool_name, tools) {
66        return clean_input.clone();
67    }
68    let serde_json::Value::Object(mut input) = clean_input.clone() else {
69        return clean_input.clone();
70    };
71    input.insert(
72        TOOL_CALL_INTENT_FIELD.into(),
73        serde_json::Value::String(intent.as_str().into()),
74    );
75    serde_json::Value::Object(input)
76}
77
78#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
79#[serde(rename_all = "snake_case")]
80pub enum MessageOrigin {
81    #[default]
82    User,
83    Watcher,
84    Interjection,
85    Internal,
86}
87
88fn is_default_origin(origin: &MessageOrigin) -> bool {
89    matches!(origin, MessageOrigin::User)
90}
91
92#[derive(Debug, Clone, Serialize, PartialEq)]
93pub struct Message {
94    pub role: MessageRole,
95    pub parts: Vec<MessagePart>,
96    pub turn_id: TurnId,
97    #[serde(default, skip_serializing_if = "is_default_origin")]
98    pub origin: MessageOrigin,
99}
100
101impl<'de> Deserialize<'de> for Message {
102    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
103    where
104        D: serde::Deserializer<'de>,
105    {
106        #[derive(Deserialize)]
107        struct RawMessage {
108            role: MessageRole,
109            parts: Vec<MessagePart>,
110            turn_id: TurnId,
111            #[serde(default)]
112            origin: MessageOrigin,
113        }
114
115        let raw = RawMessage::deserialize(deserializer)?;
116        let RawMessage {
117            role,
118            parts,
119            turn_id,
120            origin,
121        } = raw;
122        Ok(Self {
123            role,
124            parts: normalize_legacy_compact_summary(role, parts),
125            turn_id,
126            origin,
127        })
128    }
129}
130
131#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
132#[serde(rename_all = "snake_case")]
133pub enum MessageRole {
134    User,
135    Assistant,
136    System,
137    Tool,
138}
139
140#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
141#[serde(tag = "type", rename_all = "snake_case")]
142pub enum MessagePart {
143    ContextRecord(crate::context_plan::ContextRecord),
144    CompactSummary {
145        summary: String,
146        seq_start: u64,
147        seq_end: u64,
148        count: usize,
149    },
150    Text {
151        text: String,
152    },
153    Thinking {
154        thinking: String,
155        #[serde(default, skip_serializing_if = "Option::is_none")]
156        signature: Option<String>,
157    },
158    Image {
159        source: ImageSource,
160    },
161    ToolUse {
162        id: String,
163        name: String,
164        input: serde_json::Value,
165        #[serde(default, skip_serializing_if = "Option::is_none")]
166        intent: Option<ToolCallIntent>,
167    },
168    ToolResult {
169        tool_use_id: String,
170        content: String,
171        #[serde(default, skip_serializing_if = "core::ops::Not::not")]
172        is_error: bool,
173    },
174}
175
176#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
177pub struct ImageSource {
178    pub media_type: String,
179    pub data: ImageData,
180    #[serde(default, skip_serializing_if = "is_auto_image_detail")]
181    pub detail: crate::provider::ImageDetail,
182}
183
184fn is_auto_image_detail(detail: &crate::provider::ImageDetail) -> bool {
185    matches!(detail, crate::provider::ImageDetail::Auto)
186}
187
188#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
189#[serde(tag = "kind", rename_all = "snake_case")]
190pub enum ImageData {
191    Base64 {
192        data: String,
193    },
194    Path {
195        path: PathBuf,
196    },
197    Artifact {
198        id: String,
199        path: PathBuf,
200        #[serde(default, skip_serializing_if = "Option::is_none")]
201        name: Option<String>,
202    },
203}
204
205impl Message {
206    pub fn context_record(turn_id: TurnId, record: crate::context_plan::ContextRecord) -> Self {
207        Self {
208            role: MessageRole::System,
209            parts: vec![MessagePart::ContextRecord(record)],
210            turn_id,
211            origin: MessageOrigin::Internal,
212        }
213    }
214
215    pub fn user_text(turn_id: TurnId, text: impl Into<String>) -> Self {
216        Self {
217            role: MessageRole::User,
218            parts: vec![MessagePart::Text { text: text.into() }],
219            turn_id,
220            origin: MessageOrigin::User,
221        }
222    }
223
224    pub fn assistant_text(turn_id: TurnId, text: impl Into<String>) -> Self {
225        Self {
226            role: MessageRole::Assistant,
227            parts: vec![MessagePart::Text { text: text.into() }],
228            turn_id,
229            origin: MessageOrigin::User,
230        }
231    }
232
233    pub fn system_text(turn_id: TurnId, text: impl Into<String>) -> Self {
234        Self {
235            role: MessageRole::System,
236            parts: vec![MessagePart::Text { text: text.into() }],
237            turn_id,
238            origin: MessageOrigin::User,
239        }
240    }
241
242    pub fn system_compact_summary(
243        turn_id: TurnId,
244        summary: impl Into<String>,
245        seq_start: u64,
246        seq_end: u64,
247        count: usize,
248    ) -> Self {
249        Self {
250            role: MessageRole::System,
251            parts: vec![MessagePart::CompactSummary {
252                summary: summary.into(),
253                seq_start,
254                seq_end,
255                count,
256            }],
257            turn_id,
258            origin: MessageOrigin::User,
259        }
260    }
261
262    pub fn text_concat(&self) -> String {
263        let mut out = String::new();
264        for p in &self.parts {
265            match p {
266                MessagePart::ContextRecord(record) => out.push_str(&record.render_for_model()),
267                MessagePart::Text { text } => out.push_str(text),
268                MessagePart::CompactSummary { summary, .. } => out.push_str(summary),
269                _ => {}
270            }
271        }
272        out
273    }
274
275    pub fn thinking_concat(&self) -> String {
276        let mut out = String::new();
277        for p in &self.parts {
278            if let MessagePart::Thinking { thinking, .. } = p {
279                out.push_str(thinking);
280            }
281        }
282        out
283    }
284
285    pub fn thinking_signature(&self) -> Option<String> {
286        self.parts.iter().rev().find_map(|p| {
287            if let MessagePart::Thinking { signature, .. } = p {
288                signature.clone()
289            } else {
290                None
291            }
292        })
293    }
294
295    pub fn contains_context_record(&self) -> bool {
296        self.parts
297            .iter()
298            .any(|part| matches!(part, MessagePart::ContextRecord(_)))
299    }
300}
301
302pub fn retain_complete_tool_pairs(messages: &mut Vec<Message>) {
303    let use_ids: std::collections::HashSet<String> = messages
304        .iter()
305        .flat_map(|message| {
306            message.parts.iter().filter_map(|part| match part {
307                MessagePart::ToolUse { id, .. } => Some(id.clone()),
308                _ => None,
309            })
310        })
311        .collect();
312    let result_ids: std::collections::HashSet<String> = messages
313        .iter()
314        .flat_map(|message| {
315            message.parts.iter().filter_map(|part| match part {
316                MessagePart::ToolResult { tool_use_id, .. } => Some(tool_use_id.clone()),
317                _ => None,
318            })
319        })
320        .collect();
321    let mut seen_uses = std::collections::HashSet::new();
322    let mut seen_results = std::collections::HashSet::new();
323    for message in messages.iter_mut() {
324        message.parts.retain(|part| match part {
325            MessagePart::ToolUse { id, .. } => {
326                result_ids.contains(id) && seen_uses.insert(id.clone())
327            }
328            MessagePart::ToolResult { tool_use_id, .. } => {
329                use_ids.contains(tool_use_id) && seen_results.insert(tool_use_id.clone())
330            }
331            _ => true,
332        });
333    }
334    messages.retain(|message| !message.parts.is_empty());
335}
336
337pub fn normalize_tool_pairs_for_model(messages: &[Message]) -> Vec<Message> {
338    #[derive(Clone)]
339    struct ToolResultRecord {
340        part: MessagePart,
341        turn_id: TurnId,
342        origin: MessageOrigin,
343    }
344
345    let mut results = std::collections::HashMap::new();
346    for message in messages {
347        for part in &message.parts {
348            if let MessagePart::ToolResult { tool_use_id, .. } = part {
349                results
350                    .entry(tool_use_id.clone())
351                    .or_insert_with(|| ToolResultRecord {
352                        part: part.clone(),
353                        turn_id: message.turn_id.clone(),
354                        origin: message.origin,
355                    });
356            }
357        }
358    }
359
360    let mut normalized = Vec::with_capacity(messages.len() + 4);
361    for message in messages {
362        let tool_use_ids: Vec<String> = message
363            .parts
364            .iter()
365            .filter_map(|part| match part {
366                MessagePart::ToolUse { id, .. } => Some(id.clone()),
367                _ => None,
368            })
369            .collect();
370        let mut projected = message.clone();
371        projected
372            .parts
373            .retain(|part| !matches!(part, MessagePart::ToolResult { .. }));
374        if !projected.parts.is_empty() {
375            if projected.role == MessageRole::Tool {
376                projected.role = MessageRole::User;
377            }
378            normalized.push(projected);
379        }
380
381        for tool_use_id in tool_use_ids {
382            let result = results.get(&tool_use_id);
383            normalized.push(Message {
384                role: MessageRole::Tool,
385                parts: vec![result.map_or_else(
386                    || MessagePart::ToolResult {
387                        tool_use_id,
388                        content: "[tool execution interrupted — no result captured]".into(),
389                        is_error: true,
390                    },
391                    |result| result.part.clone(),
392                )],
393                turn_id: result
394                    .map(|result| result.turn_id.clone())
395                    .unwrap_or_else(|| message.turn_id.clone()),
396                origin: result.map_or(message.origin, |result| result.origin),
397            });
398        }
399    }
400    normalized
401}
402
403impl MessageRole {
404    pub fn as_str(&self) -> &'static str {
405        match self {
406            MessageRole::User => "user",
407            MessageRole::Assistant => "assistant",
408            MessageRole::System => "system",
409            MessageRole::Tool => "tool",
410        }
411    }
412}
413
414fn normalize_legacy_compact_summary(
415    role: MessageRole,
416    parts: Vec<MessagePart>,
417) -> Vec<MessagePart> {
418    if role != MessageRole::System {
419        return parts;
420    }
421    if parts.len() != 1 {
422        return parts;
423    }
424    let MessagePart::Text { text } = &parts[0] else {
425        return parts;
426    };
427    let Some((summary, seq_start, seq_end, count)) = parse_legacy_compact_summary_text(text) else {
428        return parts;
429    };
430    vec![MessagePart::CompactSummary {
431        summary,
432        seq_start,
433        seq_end,
434        count,
435    }]
436}
437
438pub(crate) fn parse_legacy_compact_summary_text(text: &str) -> Option<(String, u64, u64, usize)> {
439    let start_marker = "[atman:compact ";
440    let start = text.rfind(start_marker)?;
441    let after = &text[start + start_marker.len()..];
442    let end = after.find(']')?;
443    let inner = &after[..end];
444    let mut seq_start = None;
445    let mut seq_end = None;
446    let mut count = None;
447    for token in inner.split_whitespace() {
448        let Some((k, v)) = token.split_once('=') else {
449            continue;
450        };
451        match k {
452            "seq_start" => seq_start = v.parse().ok(),
453            "seq_end" => seq_end = v.parse().ok(),
454            "count" => count = v.parse().ok(),
455            _ => {}
456        }
457    }
458    let summary = text[..start].trim_end().to_string();
459    Some((summary, seq_start?, seq_end?, count?))
460}
461
462#[cfg(test)]
463mod tests {
464    use super::*;
465
466    #[test]
467    fn user_text_roundtrips_via_serde_json() {
468        let msg = Message::user_text(TurnId::now(), "hello");
469        let s = serde_json::to_string(&msg).unwrap();
470        let back: Message = serde_json::from_str(&s).unwrap();
471        assert_eq!(msg, back);
472    }
473
474    #[test]
475    fn legacy_compact_summary_deserializes_to_structured_variant() {
476        let turn_id = TurnId::now();
477        let msg = Message {
478            role: MessageRole::System,
479            parts: vec![MessagePart::Text {
480                text: "handoff\n\n[atman:compact seq_start=2 seq_end=7 count=6]".into(),
481            }],
482            turn_id,
483            origin: MessageOrigin::User,
484        };
485        let s = serde_json::to_string(&msg).unwrap();
486        let back: Message = serde_json::from_str(&s).unwrap();
487        assert!(matches!(
488            back.parts.as_slice(),
489            [MessagePart::CompactSummary { .. }]
490        ));
491        assert_eq!(back.text_concat(), "handoff");
492    }
493
494    #[test]
495    fn text_concat_skips_non_text_parts() {
496        let msg = Message {
497            role: MessageRole::User,
498            parts: vec![
499                MessagePart::Text { text: "a ".into() },
500                MessagePart::Image {
501                    source: ImageSource {
502                        media_type: "image/png".into(),
503                        data: ImageData::Path {
504                            path: PathBuf::from("/tmp/x.png"),
505                        },
506                        detail: crate::provider::ImageDetail::Auto,
507                    },
508                },
509                MessagePart::Text { text: "b".into() },
510            ],
511            turn_id: TurnId::now(),
512            origin: MessageOrigin::User,
513        };
514        assert_eq!(msg.text_concat(), "a b");
515    }
516
517    #[test]
518    fn tool_result_is_error_defaults_to_false_and_skips_serialize_when_false() {
519        let msg = Message {
520            role: MessageRole::Tool,
521            parts: vec![MessagePart::ToolResult {
522                tool_use_id: "toolu_1".into(),
523                content: "ok".into(),
524                is_error: false,
525            }],
526            turn_id: TurnId::now(),
527            origin: MessageOrigin::User,
528        };
529        let s = serde_json::to_string(&msg).unwrap();
530        assert!(!s.contains("is_error"), "should skip when false: {s}");
531
532        let err_msg = Message {
533            role: MessageRole::Tool,
534            parts: vec![MessagePart::ToolResult {
535                tool_use_id: "toolu_1".into(),
536                content: "nope".into(),
537                is_error: true,
538            }],
539            turn_id: TurnId::now(),
540            origin: MessageOrigin::User,
541        };
542        let s = serde_json::to_string(&err_msg).unwrap();
543        assert!(s.contains("\"is_error\":true"), "{s}");
544    }
545
546    #[test]
547    fn role_as_str_matches_wire_format() {
548        assert_eq!(MessageRole::User.as_str(), "user");
549        assert_eq!(MessageRole::Assistant.as_str(), "assistant");
550        assert_eq!(MessageRole::System.as_str(), "system");
551        assert_eq!(MessageRole::Tool.as_str(), "tool");
552    }
553
554    #[test]
555    fn default_origin_is_user() {
556        assert_eq!(MessageOrigin::default(), MessageOrigin::User);
557    }
558
559    #[test]
560    fn user_origin_skipped_in_json() {
561        let msg = Message::user_text(TurnId::now(), "hi");
562        let s = serde_json::to_string(&msg).unwrap();
563        assert!(
564            !s.contains("origin"),
565            "default origin should not be serialized: {s}"
566        );
567    }
568
569    #[test]
570    fn watcher_origin_serialized() {
571        let mut msg = Message::user_text(TurnId::now(), "watcher event");
572        msg.origin = MessageOrigin::Watcher;
573        let s = serde_json::to_string(&msg).unwrap();
574        assert!(s.contains("\"origin\":\"watcher\""), "{s}");
575        let back: Message = serde_json::from_str(&s).unwrap();
576        assert_eq!(back.origin, MessageOrigin::Watcher);
577    }
578
579    #[test]
580    fn old_json_without_origin_defaults_to_user() {
581        let json = r#"{"role":"user","parts":[{"type":"text","text":"legacy"}],"turn_id":"019f0000-0000-7000-0000-000000000001"}"#;
582        let msg: Message = serde_json::from_str(json).unwrap();
583        assert_eq!(msg.origin, MessageOrigin::User);
584        assert_eq!(msg.text_concat(), "legacy");
585    }
586
587    #[test]
588    fn context_record_message_round_trips_as_internal_system_context() {
589        let message = Message::context_record(
590            TurnId::now(),
591            crate::context_plan::ContextRecord::new(
592                "session.workspace",
593                1,
594                crate::context_plan::ContextRecordAuthority::Runtime,
595                crate::context_plan::ContextRecordRetention::Latest,
596                crate::context_plan::ContextRecordBody::text("/workspace"),
597            ),
598        );
599        let encoded = serde_json::to_string(&message).unwrap();
600        let decoded: Message = serde_json::from_str(&encoded).unwrap();
601
602        assert_eq!(decoded, message);
603        assert_eq!(decoded.role, MessageRole::System);
604        assert_eq!(decoded.origin, MessageOrigin::Internal);
605        assert!(decoded.text_concat().contains("/workspace"));
606    }
607
608    #[test]
609    fn tool_call_intent_normalizes_whitespace_and_caps_unicode_chars() {
610        let raw = format!("  inspect\n\t{}  ", "界".repeat(200));
611        let intent = ToolCallIntent::new(raw).unwrap();
612        assert_eq!(intent.as_str().chars().count(), TOOL_CALL_INTENT_MAX_CHARS);
613        assert!(intent.as_str().starts_with("inspect 界"));
614        assert!(!intent.as_str().contains('\n'));
615    }
616
617    #[test]
618    fn legacy_tool_use_defaults_intent_to_none() {
619        let part: MessagePart = serde_json::from_value(serde_json::json!({
620            "type": "tool_use",
621            "id": "call-1",
622            "name": "probe",
623            "input": {"value": 1}
624        }))
625        .unwrap();
626        assert!(matches!(part, MessagePart::ToolUse { intent: None, .. }));
627    }
628
629    #[test]
630    fn model_normalization_splits_parallel_results_in_call_order() {
631        let turn = TurnId::now();
632        let messages = vec![
633            Message {
634                role: MessageRole::Assistant,
635                parts: vec![
636                    MessagePart::ToolUse {
637                        id: "call-b".into(),
638                        name: "probe".into(),
639                        input: serde_json::json!({}),
640                        intent: None,
641                    },
642                    MessagePart::ToolUse {
643                        id: "call-a".into(),
644                        name: "probe".into(),
645                        input: serde_json::json!({}),
646                        intent: None,
647                    },
648                ],
649                turn_id: turn.clone(),
650                origin: MessageOrigin::User,
651            },
652            Message {
653                role: MessageRole::Tool,
654                parts: vec![
655                    MessagePart::ToolResult {
656                        tool_use_id: "call-a".into(),
657                        content: "A".into(),
658                        is_error: false,
659                    },
660                    MessagePart::ToolResult {
661                        tool_use_id: "call-b".into(),
662                        content: "B".into(),
663                        is_error: false,
664                    },
665                ],
666                turn_id: turn,
667                origin: MessageOrigin::User,
668            },
669        ];
670
671        let normalized = normalize_tool_pairs_for_model(&messages);
672        let results: Vec<(&str, &str)> = normalized
673            .iter()
674            .flat_map(|message| &message.parts)
675            .filter_map(|part| match part {
676                MessagePart::ToolResult {
677                    tool_use_id,
678                    content,
679                    ..
680                } => Some((tool_use_id.as_str(), content.as_str())),
681                _ => None,
682            })
683            .collect();
684        assert_eq!(results, vec![("call-b", "B"), ("call-a", "A")]);
685        assert_eq!(normalize_tool_pairs_for_model(&normalized), normalized);
686    }
687
688    #[test]
689    fn model_normalization_preserves_mixed_content_and_drops_orphan_results() {
690        let turn = TurnId::now();
691        let messages = vec![
692            Message {
693                role: MessageRole::Assistant,
694                parts: vec![
695                    MessagePart::Text {
696                        text: "checking".into(),
697                    },
698                    MessagePart::ToolUse {
699                        id: "call-ok".into(),
700                        name: "probe".into(),
701                        input: serde_json::json!({}),
702                        intent: None,
703                    },
704                ],
705                turn_id: turn.clone(),
706                origin: MessageOrigin::User,
707            },
708            Message {
709                role: MessageRole::Tool,
710                parts: vec![
711                    MessagePart::Text {
712                        text: "preserve me".into(),
713                    },
714                    MessagePart::ToolResult {
715                        tool_use_id: "call-ok".into(),
716                        content: "done".into(),
717                        is_error: false,
718                    },
719                    MessagePart::ToolResult {
720                        tool_use_id: "orphan".into(),
721                        content: "drop me".into(),
722                        is_error: false,
723                    },
724                ],
725                turn_id: turn,
726                origin: MessageOrigin::Watcher,
727            },
728            Message {
729                role: MessageRole::User,
730                parts: Vec::new(),
731                turn_id: TurnId::now(),
732                origin: MessageOrigin::User,
733            },
734        ];
735
736        let normalized = normalize_tool_pairs_for_model(&messages);
737        assert_eq!(normalized.len(), 3);
738        assert_eq!(normalized[0].text_concat(), "checking");
739        assert!(matches!(
740            &normalized[1].parts[..],
741            [MessagePart::ToolResult { tool_use_id, content, .. }]
742                if tool_use_id == "call-ok" && content == "done"
743        ));
744        assert_eq!(normalized[2].role, MessageRole::User);
745        assert_eq!(normalized[2].origin, MessageOrigin::Watcher);
746        assert_eq!(normalized[2].text_concat(), "preserve me");
747        assert!(!normalized.iter().any(|message| {
748            message.parts.iter().any(|part| {
749                matches!(part, MessagePart::ToolResult { tool_use_id, .. } if tool_use_id == "orphan")
750            })
751        }));
752    }
753}