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