ai_agent/utils/
messages.rs

1// Source: /data/home/swei/claudecode/openclaudecode/src/constants/messages.ts
2//! Message utilities and helpers
3//! Translated from /data/home/swei/claudecode/openclaudecode/src/utils/messages.ts
4
5use std::collections::{HashMap, HashSet};
6
7use serde::{Deserialize, Serialize};
8
9/// Message types
10#[derive(Debug, Clone, Serialize, Deserialize)]
11#[serde(tag = "type")]
12pub enum Message {
13    User(UserMessage),
14    Assistant(AssistantMessage),
15    Progress(ProgressMessage),
16    Attachment(AttachmentMessage),
17    System(SystemMessage),
18}
19
20/// User message
21#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct UserMessage {
23    pub message: MessageContent,
24    #[serde(skip_serializing_if = "Option::is_none")]
25    pub is_meta: Option<bool>,
26    #[serde(skip_serializing_if = "Option::is_none")]
27    pub is_visible_in_transcript_only: Option<bool>,
28    #[serde(skip_serializing_if = "Option::is_none")]
29    pub is_virtual: Option<bool>,
30    #[serde(skip_serializing_if = "Option::is_none")]
31    pub is_compact_summary: Option<bool>,
32    #[serde(skip_serializing_if = "Option::is_none")]
33    pub summarize_metadata: Option<SummarizeMetadata>,
34    #[serde(skip_serializing_if = "Option::is_none")]
35    pub tool_use_result: Option<serde_json::Value>,
36    #[serde(skip_serializing_if = "Option::is_none")]
37    pub mcp_meta: Option<serde_json::Value>,
38    pub uuid: String,
39    pub timestamp: String,
40    #[serde(skip_serializing_if = "Option::is_none")]
41    pub image_paste_ids: Option<Vec<u32>>,
42    #[serde(skip_serializing_if = "Option::is_none")]
43    pub source_tool_assistant_uuid: Option<String>,
44    #[serde(skip_serializing_if = "Option::is_none")]
45    pub permission_mode: Option<String>,
46    #[serde(skip_serializing_if = "Option::is_none")]
47    pub origin: Option<MessageOrigin>,
48}
49
50/// Summarize metadata
51#[derive(Debug, Clone, Serialize, Deserialize)]
52pub struct SummarizeMetadata {
53    pub messages_summarized: u32,
54    #[serde(skip_serializing_if = "Option::is_none")]
55    pub user_context: Option<String>,
56    #[serde(skip_serializing_if = "Option::is_none")]
57    pub direction: Option<String>,
58}
59
60/// Message origin
61#[derive(Debug, Clone, Serialize, Deserialize)]
62pub struct MessageOrigin {
63    #[serde(rename = "type")]
64    pub origin_type: String,
65}
66
67/// Message content (can be string or array of content blocks)
68#[derive(Debug, Clone, Serialize, Deserialize)]
69#[serde(untagged)]
70pub enum MessageContent {
71    String(String),
72    Blocks(Vec<ContentBlock>),
73}
74
75/// Content block
76#[derive(Debug, Clone, Serialize, Deserialize)]
77#[serde(tag = "type", rename_all = "snake_case")]
78pub enum ContentBlock {
79    Text {
80        text: String,
81    },
82    Image {
83        source: ImageSource,
84    },
85    ToolUse {
86        id: String,
87        name: String,
88        input: serde_json::Value,
89    },
90    ToolResult {
91        tool_use_id: String,
92        content: Option<Vec<ContentBlock>>,
93        #[serde(skip_serializing_if = "Option::is_none")]
94        is_error: Option<bool>,
95    },
96    // Server-side tool results
97    #[serde(rename = "server_tool_use")]
98    ServerToolUse {
99        id: String,
100        name: String,
101        input: serde_json::Value,
102    },
103    #[serde(rename = "mcp_tool_use")]
104    McpToolUse {
105        id: String,
106        name: String,
107        input: serde_json::Value,
108    },
109    #[serde(rename = "advisor_tool_result")]
110    AdvisorToolResult {
111        tool_use_id: String,
112        content: serde_json::Value,
113    },
114    #[serde(rename = "web_search_tool_result")]
115    WebSearchToolResult {
116        tool_use_id: String,
117        content: serde_json::Value,
118    },
119    #[serde(rename = "web_fetch_tool_result")]
120    WebFetchToolResult {
121        tool_use_id: String,
122        content: serde_json::Value,
123    },
124    #[serde(rename = "tool_reference")]
125    ToolReference {
126        tool_name: String,
127    },
128}
129
130/// Image source
131#[derive(Debug, Clone, Serialize, Deserialize)]
132pub struct ImageSource {
133    #[serde(rename = "type")]
134    pub source_type: String,
135    pub media_type: String,
136    pub data: String,
137}
138
139/// Assistant message
140#[derive(Debug, Clone, Serialize, Deserialize)]
141pub struct AssistantMessage {
142    pub message: AssistantMessageContent,
143    #[serde(skip_serializing_if = "Option::is_none")]
144    pub request_id: Option<String>,
145    #[serde(skip_serializing_if = "Option::is_none")]
146    pub api_error: Option<serde_json::Value>,
147    #[serde(skip_serializing_if = "Option::is_none")]
148    pub error: Option<serde_json::Value>,
149    #[serde(skip_serializing_if = "Option::is_none")]
150    pub error_details: Option<String>,
151    #[serde(skip_serializing_if = "Option::is_none")]
152    pub is_api_error_message: Option<bool>,
153    #[serde(skip_serializing_if = "Option::is_none")]
154    pub is_virtual: Option<bool>,
155    #[serde(skip_serializing_if = "Option::is_none")]
156    pub is_meta: Option<bool>,
157    #[serde(skip_serializing_if = "Option::is_none")]
158    pub advisor_model: Option<String>,
159    pub uuid: String,
160    pub timestamp: String,
161    #[serde(skip_serializing_if = "Option::is_none")]
162    pub parent_uuid: Option<String>,
163}
164
165/// Assistant message content
166#[derive(Debug, Clone, Serialize, Deserialize)]
167pub struct AssistantMessageContent {
168    pub id: String,
169    #[serde(skip_serializing_if = "Option::is_none")]
170    pub container: Option<String>,
171    pub model: String,
172    pub role: String,
173    #[serde(skip_serializing_if = "Option::is_none")]
174    pub stop_reason: Option<String>,
175    #[serde(skip_serializing_if = "Option::is_none")]
176    pub stop_sequence: Option<String>,
177    #[serde(rename = "type")]
178    pub message_type: String,
179    pub usage: Option<Usage>,
180    pub content: Vec<serde_json::Value>,
181    #[serde(skip_serializing_if = "Option::is_none")]
182    pub context_management: Option<serde_json::Value>,
183}
184
185/// Usage statistics
186#[derive(Debug, Clone, Default, Serialize, Deserialize)]
187pub struct Usage {
188    #[serde(rename = "input_tokens")]
189    pub input_tokens: u32,
190    #[serde(rename = "output_tokens")]
191    pub output_tokens: u32,
192    #[serde(rename = "cache_creation_input_tokens")]
193    pub cache_creation_input_tokens: u32,
194    #[serde(rename = "cache_read_input_tokens")]
195    pub cache_read_input_tokens: u32,
196    #[serde(rename = "server_tool_use")]
197    pub server_tool_use: ServerToolUse,
198    #[serde(skip_serializing_if = "Option::is_none")]
199    pub service_tier: Option<String>,
200    #[serde(skip_serializing_if = "Option::is_none")]
201    pub cache_creation: Option<CacheCreation>,
202    #[serde(skip_serializing_if = "Option::is_none")]
203    pub inference_geo: Option<String>,
204    #[serde(skip_serializing_if = "Option::is_none")]
205    pub iterations: Option<u32>,
206    #[serde(skip_serializing_if = "Option::is_none")]
207    pub speed: Option<f64>,
208}
209
210/// Server tool use stats
211#[derive(Debug, Clone, Default, Serialize, Deserialize)]
212pub struct ServerToolUse {
213    #[serde(rename = "web_search_requests")]
214    pub web_search_requests: u32,
215    #[serde(rename = "web_fetch_requests")]
216    pub web_fetch_requests: u32,
217}
218
219/// Cache creation stats
220#[derive(Debug, Clone, Default, Serialize, Deserialize)]
221pub struct CacheCreation {
222    #[serde(rename = "ephemeral_1h_input_tokens")]
223    pub ephemeral_1h_input_tokens: u32,
224    #[serde(rename = "ephemeral_5m_input_tokens")]
225    pub ephemeral_5m_input_tokens: u32,
226}
227
228/// Progress message
229#[derive(Debug, Clone, Serialize, Deserialize)]
230pub struct ProgressMessage<T = serde_json::Value> {
231    #[serde(rename = "type")]
232    pub data_type: String,
233    pub data: T,
234    #[serde(skip_serializing_if = "Option::is_none")]
235    pub tool_use_id: Option<String>,
236    #[serde(skip_serializing_if = "Option::is_none")]
237    pub parent_tool_use_id: Option<String>,
238    pub uuid: String,
239    pub timestamp: String,
240    #[serde(skip_serializing_if = "Option::is_none")]
241    pub parent_uuid: Option<String>,
242}
243
244/// Attachment message
245#[derive(Debug, Clone, Serialize, Deserialize)]
246pub struct AttachmentMessage {
247    pub attachment: serde_json::Value,
248    pub uuid: String,
249    pub timestamp: String,
250    #[serde(skip_serializing_if = "Option::is_none")]
251    pub parent_uuid: Option<String>,
252}
253
254/// System message
255#[derive(Debug, Clone, Serialize, Deserialize)]
256pub struct SystemMessage {
257    pub message: SystemMessageContent,
258    pub uuid: String,
259    pub timestamp: String,
260    #[serde(skip_serializing_if = "Option::is_none")]
261    pub parent_uuid: Option<String>,
262}
263
264/// System message content
265#[derive(Debug, Clone, Serialize, Deserialize)]
266pub struct SystemMessageContent {
267    #[serde(rename = "type")]
268    pub message_type: String,
269    #[serde(skip_serializing_if = "Option::is_none")]
270    pub subtype: Option<String>,
271    pub content: String,
272    #[serde(skip_serializing_if = "Option::is_none")]
273    pub level: Option<SystemMessageLevel>,
274}
275
276/// System message level
277#[derive(Debug, Clone, Serialize, Deserialize)]
278#[serde(rename_all = "lowercase")]
279pub enum SystemMessageLevel {
280    Info,
281    Warning,
282    Error,
283}
284
285// === Constants ===
286
287pub const INTERRUPT_MESSAGE: &str = "[Request interrupted by user]";
288pub const INTERRUPT_MESSAGE_FOR_TOOL_USE: &str = "[Request interrupted by user for tool use]";
289pub const CANCEL_MESSAGE: &str = "The user doesn't want to take this action right now. STOP what you are doing and wait for the user to tell you how to proceed.";
290pub const REJECT_MESSAGE: &str = "The user doesn't want to proceed with this tool use. The tool use was rejected (eg. if it was a file edit, the new_string was NOT written to the file). STOP what you are doing and wait for the user to tell you how to proceed.";
291pub const REJECT_MESSAGE_WITH_REASON_PREFIX: &str = "The user doesn't want to proceed with this tool use. The tool use was rejected (eg. if it was a file edit, the new_string was NOT written to the file). To tell you how to proceed, the user said:\n";
292pub const SUBAGENT_REJECT_MESSAGE: &str = "Permission for this tool use was denied. The tool use was rejected (eg. if it was a file edit, the new_string was NOT written to the file). Try a different approach or report the limitation to complete your task.";
293pub const SUBAGENT_REJECT_MESSAGE_WITH_REASON_PREFIX: &str = "Permission for this tool use was denied. The tool use was rejected (eg. if it was a file edit, the new_string was NOT written to the file). The user said:\n";
294pub const NO_RESPONSE_REQUESTED: &str = "No response requested.";
295pub const SYNTHETIC_MODEL: &str = "<synthetic>";
296
297// Denial workaround guidance
298pub const DENIAL_WORKAROUND_GUIDANCE: &str = "IMPORTANT: You *may* attempt to accomplish this action using other tools that might naturally be used to accomplish this goal, e.g. using head instead of cat. But you *should not* attempt to work around this denial in malicious ways, e.g. do not use your ability to run tests to execute non-test actions. You should only try to work around this restriction in reasonable ways that do not attempt to bypass the intent behind this denial. If you believe this capability is essential to complete your request, STOP and explain to the user what you were trying to do and why you need this permission. Let the user decide how to proceed.";
299
300const AUTO_MODE_REJECTION_PREFIX: &str = "Permission for this action has been denied. Reason: ";
301
302// === Functions ===
303
304/// Build rejection message for auto mode classifier denials
305pub fn build_yolo_rejection_message(reason: &str) -> String {
306    format!(
307        "{}{}. If you have other tasks that don't depend on this action, continue working on those. {} To allow this type of action in the future, the user can add a Bash permission rule to their settings.",
308        AUTO_MODE_REJECTION_PREFIX, reason, DENIAL_WORKAROUND_GUIDANCE
309    )
310}
311
312/// Build message for when classifier is unavailable
313pub fn build_classifier_unavailable_message(tool_name: &str, classifier_model: &str) -> String {
314    format!(
315        "{} is temporarily unavailable, so auto mode cannot determine the safety of {} right now. Wait briefly and then try this action again. If it keeps failing, continue with other tasks that don't require this action and come back to it later. Note: reading files, searching code, and other read-only operations do not require the classifier and can still be used.",
316        classifier_model, tool_name
317    )
318}
319
320/// Check if tool result message is a classifier denial
321pub fn is_classifier_denial(content: &str) -> bool {
322    content.starts_with(AUTO_MODE_REJECTION_PREFIX)
323}
324
325/// Auto reject message
326pub fn auto_reject_message(tool_name: &str) -> String {
327    format!(
328        "Permission to use {} has been denied. {}",
329        tool_name, DENIAL_WORKAROUND_GUIDANCE
330    )
331}
332
333/// Don't ask reject message
334pub fn dont_ask_reject_message(tool_name: &str) -> String {
335    format!(
336        "Permission to use {} has been denied because Claude Code is running in don't ask mode. {}",
337        tool_name, DENIAL_WORKAROUND_GUIDANCE
338    )
339}
340
341/// Derive short message ID (6-char base36) from UUID
342pub fn derive_short_message_id(uuid: &str) -> String {
343    // Take first 10 hex chars from the UUID (skipping dashes)
344    let hex: String = uuid.replace('-', "").chars().take(10).collect();
345    // Convert to base36 for shorter representation, take 6 chars
346    let parsed = u64::from_str_radix(&hex, 16).unwrap_or(0);
347    let base36 = format!("{:36}", parsed);
348    base36.chars().take(6).collect()
349}
350
351/// Create a user message
352pub fn create_user_message(content: impl Into<MessageContent>) -> UserMessage {
353    let content = content.into();
354    UserMessage {
355        message: content,
356        is_meta: None,
357        is_visible_in_transcript_only: None,
358        is_virtual: None,
359        is_compact_summary: None,
360        summarize_metadata: None,
361        tool_use_result: None,
362        mcp_meta: None,
363        uuid: uuid::Uuid::new_v4().to_string(),
364        timestamp: chrono::Utc::now().to_rfc3339(),
365        image_paste_ids: None,
366        source_tool_assistant_uuid: None,
367        permission_mode: None,
368        origin: None,
369    }
370}
371
372/// Create assistant message
373pub fn create_assistant_message(content: Vec<serde_json::Value>) -> AssistantMessage {
374    AssistantMessage {
375        message: AssistantMessageContent {
376            id: uuid::Uuid::new_v4().to_string(),
377            container: None,
378            model: SYNTHETIC_MODEL.to_string(),
379            role: "assistant".to_string(),
380            stop_reason: Some("stop_sequence".to_string()),
381            stop_sequence: Some("".to_string()),
382            message_type: "message".to_string(),
383            usage: Some(Usage::default()),
384            content,
385            context_management: None,
386        },
387        request_id: None,
388        api_error: None,
389        error: None,
390        error_details: None,
391        is_api_error_message: Some(false),
392        is_virtual: None,
393        is_meta: None,
394        advisor_model: None,
395        uuid: uuid::Uuid::new_v4().to_string(),
396        timestamp: chrono::Utc::now().to_rfc3339(),
397        parent_uuid: None,
398    }
399}
400
401/// Create progress message
402pub fn create_progress_message(
403    tool_use_id: &str,
404    parent_tool_use_id: &str,
405    data: serde_json::Value,
406) -> ProgressMessage {
407    ProgressMessage {
408        data_type: "progress".to_string(),
409        data,
410        tool_use_id: Some(tool_use_id.to_string()),
411        parent_tool_use_id: Some(parent_tool_use_id.to_string()),
412        uuid: uuid::Uuid::new_v4().to_string(),
413        timestamp: chrono::Utc::now().to_rfc3339(),
414        parent_uuid: None,
415    }
416}
417
418/// Create tool result stop message
419pub fn create_tool_result_stop_message(tool_use_id: &str) -> serde_json::Value {
420    serde_json::json!({
421        "type": "tool_result",
422        "content": CANCEL_MESSAGE,
423        "is_error": true,
424        "tool_use_id": tool_use_id
425    })
426}
427
428/// XML tags for command input/output
429pub const COMMAND_MESSAGE_TAG: &str = "command-message";
430pub const COMMAND_NAME_TAG: &str = "command-name";
431
432/// Create a synthetic user caveat message (informs the model the user typed something)
433pub fn create_synthetic_user_caveat_message() -> Message {
434    let content = "The user didn't say anything. Continue working.".to_string();
435    Message::User(UserMessage {
436        message: MessageContent::String(content),
437        is_meta: Some(true),
438        is_visible_in_transcript_only: None,
439        is_virtual: None,
440        is_compact_summary: None,
441        summarize_metadata: None,
442        tool_use_result: None,
443        mcp_meta: None,
444        uuid: uuid::Uuid::new_v4().to_string(),
445        timestamp: chrono::Utc::now().to_rfc3339(),
446        image_paste_ids: None,
447        source_tool_assistant_uuid: None,
448        permission_mode: None,
449        origin: None,
450    })
451}
452
453/// Create a system message
454pub fn create_system_message(content: impl Into<String>, level: SystemMessageLevel) -> Message {
455    Message::System(SystemMessage {
456        message: SystemMessageContent {
457            message_type: "system".to_string(),
458            subtype: None,
459            content: content.into(),
460            level: Some(level),
461        },
462        uuid: uuid::Uuid::new_v4().to_string(),
463        timestamp: chrono::Utc::now().to_rfc3339(),
464        parent_uuid: None,
465    })
466}
467
468/// Create a system local command message (for command output that's visible but not sent to model)
469pub fn create_system_local_command_message(content: impl Into<String>) -> Message {
470    Message::System(SystemMessage {
471        message: SystemMessageContent {
472            message_type: "system".to_string(),
473            subtype: Some("local_command".to_string()),
474            content: content.into(),
475            level: None,
476        },
477        uuid: uuid::Uuid::new_v4().to_string(),
478        timestamp: chrono::Utc::now().to_rfc3339(),
479        parent_uuid: None,
480    })
481}
482
483/// Create a user interruption message
484pub fn create_user_interruption_message(tool_use: bool) -> Message {
485    let content = if tool_use {
486        INTERRUPT_MESSAGE_FOR_TOOL_USE.to_string()
487    } else {
488        INTERRUPT_MESSAGE.to_string()
489    };
490    Message::User(UserMessage {
491        message: MessageContent::String(content),
492        is_meta: None,
493        is_visible_in_transcript_only: None,
494        is_virtual: None,
495        is_compact_summary: None,
496        summarize_metadata: None,
497        tool_use_result: None,
498        mcp_meta: None,
499        uuid: uuid::Uuid::new_v4().to_string(),
500        timestamp: chrono::Utc::now().to_rfc3339(),
501        image_paste_ids: None,
502        source_tool_assistant_uuid: None,
503        permission_mode: None,
504        origin: None,
505    })
506}
507
508/// Format command input with XML tags: `<command-message>name</command-message>\n<command-name>/name</command-name>`
509pub fn format_command_input_tags(command_name: &str, args: &str) -> String {
510    let mut parts = vec![
511        format!("<{COMMAND_MESSAGE_TAG}>{command_name}</{COMMAND_MESSAGE_TAG}>"),
512        format!("<{COMMAND_NAME_TAG}>/{command_name}</{COMMAND_NAME_TAG}>"),
513    ];
514    if !args.trim().is_empty() {
515        parts.push(format!("<command-args>{args}</command-args>"));
516    }
517    parts.join("\n")
518}
519
520/// Check if a message is a system local command message
521pub fn is_system_local_command_message(message: &Message) -> bool {
522    match message {
523        Message::System(sys) => sys.message.subtype.as_deref() == Some("local_command"),
524        _ => false,
525    }
526}
527
528/// Check if a message is a compact boundary message (used by /compact)
529pub fn is_compact_boundary_message(message: &Message) -> bool {
530    match message {
531        Message::User(user) => user.is_compact_summary == Some(true),
532        Message::System(sys) => sys.message.subtype.as_deref() == Some("compact"),
533        _ => false,
534    }
535}
536
537/// Extract tag from HTML-like content
538pub fn extract_tag(html: &str, tag_name: &str) -> Option<String> {
539    use regex::Regex;
540
541    if html.trim().is_empty() || tag_name.trim().is_empty() {
542        return None;
543    }
544
545    let escaped_tag = tag_name.replace(
546        [
547            '.', '*', '+', '?', '^', '$', '{', '}', '[', ']', '(', ')', '|', '\\',
548        ],
549        "\\$&",
550    );
551
552    let pattern = format!(
553        r"<{}(?:\s+[^>]*)?>([\s\S]*?)</{}>",
554        escaped_tag, escaped_tag
555    );
556
557    let re = Regex::new(&pattern).ok()?;
558
559    let mut depth = 0i32;
560    let mut last_index = 0;
561
562    let opening_tag_re = Regex::new(&format!(r"<{}(?:\s+[^>]*)?>", escaped_tag)).ok()?;
563    let closing_tag_re = Regex::new(&format!(r"</{}>", escaped_tag)).ok()?;
564
565    for caps in re.captures_iter(html) {
566        let content = caps.get(1)?.as_str();
567        let start = caps.get(0)?.start();
568
569        depth = 0;
570
571        for _ in opening_tag_re.find_iter(&html[..start]) {
572            depth += 1;
573        }
574
575        for _ in closing_tag_re.find_iter(&html[..start]) {
576            depth -= 1;
577        }
578
579        if depth == 0 && !content.is_empty() {
580            return Some(content.to_string());
581        }
582
583        last_index = start + caps.get(0)?.len();
584    }
585
586    None
587}
588
589/// Check if message is not empty
590pub fn is_not_empty_message(message: &Message) -> bool {
591    match message {
592        Message::Progress(_) | Message::Attachment(_) | Message::System(_) => true,
593        Message::User(user) => {
594            match &user.message {
595                MessageContent::String(s) => !s.trim().is_empty(),
596                MessageContent::Blocks(blocks) => {
597                    if blocks.is_empty() {
598                        return false;
599                    }
600                    // Skip multi-block messages for now
601                    if blocks.len() > 1 {
602                        return true;
603                    }
604                    // Check first block
605                    match &blocks[0] {
606                        ContentBlock::Text { text } => {
607                            !text.trim().is_empty()
608                                && text != NO_RESPONSE_REQUESTED
609                                && text != INTERRUPT_MESSAGE_FOR_TOOL_USE
610                        }
611                        _ => true,
612                    }
613                }
614            }
615        }
616        Message::Assistant(assistant) => !assistant.message.content.is_empty(),
617    }
618}
619
620/// Normalized message types
621#[derive(Debug, Clone, Serialize, Deserialize)]
622#[serde(tag = "type")]
623pub enum NormalizedMessage {
624    User(NormalizedUserMessage),
625    Assistant(NormalizedAssistantMessage),
626    Progress(ProgressMessage),
627    Attachment(AttachmentMessage),
628    System(SystemMessage),
629}
630
631/// Normalized user message
632#[derive(Debug, Clone, Serialize, Deserialize)]
633pub struct NormalizedUserMessage {
634    pub message: MessageContent,
635    #[serde(flatten)]
636    pub extra: UserMessageExtra,
637}
638
639/// Normalized assistant message
640#[derive(Debug, Clone, Serialize, Deserialize)]
641pub struct NormalizedAssistantMessage {
642    pub message: AssistantMessageContent,
643    #[serde(flatten)]
644    pub extra: AssistantMessageExtra,
645}
646
647/// Extra fields for normalized user message
648#[derive(Debug, Clone, Default, Serialize, Deserialize)]
649pub struct UserMessageExtra {
650    #[serde(skip_serializing_if = "Option::is_none")]
651    pub is_meta: Option<bool>,
652    #[serde(skip_serializing_if = "Option::is_none")]
653    pub is_visible_in_transcript_only: Option<bool>,
654    #[serde(skip_serializing_if = "Option::is_none")]
655    pub is_virtual: Option<bool>,
656    #[serde(skip_serializing_if = "Option::is_none")]
657    pub is_compact_summary: Option<bool>,
658    #[serde(skip_serializing_if = "Option::is_none")]
659    pub summarize_metadata: Option<serde_json::Value>,
660    #[serde(skip_serializing_if = "Option::is_none")]
661    pub tool_use_result: Option<serde_json::Value>,
662    #[serde(skip_serializing_if = "Option::is_none")]
663    pub mcp_meta: Option<serde_json::Value>,
664    pub uuid: String,
665    pub timestamp: String,
666    #[serde(skip_serializing_if = "Option::is_none")]
667    pub image_paste_ids: Option<Vec<u32>>,
668    #[serde(skip_serializing_if = "Option::is_none")]
669    pub source_tool_assistant_uuid: Option<String>,
670    #[serde(skip_serializing_if = "Option::is_none")]
671    pub permission_mode: Option<String>,
672    #[serde(skip_serializing_if = "Option::is_none")]
673    pub origin: Option<MessageOrigin>,
674    #[serde(skip_serializing_if = "Option::is_none")]
675    pub parent_uuid: Option<String>,
676}
677
678/// Extra fields for normalized assistant message
679#[derive(Debug, Clone, Default, Serialize, Deserialize)]
680pub struct AssistantMessageExtra {
681    #[serde(skip_serializing_if = "Option::is_none")]
682    pub request_id: Option<String>,
683    #[serde(skip_serializing_if = "Option::is_none")]
684    pub api_error: Option<serde_json::Value>,
685    #[serde(skip_serializing_if = "Option::is_none")]
686    pub error: Option<serde_json::Value>,
687    #[serde(skip_serializing_if = "Option::is_none")]
688    pub error_details: Option<String>,
689    #[serde(skip_serializing_if = "Option::is_none")]
690    pub is_api_error_message: Option<bool>,
691    #[serde(skip_serializing_if = "Option::is_none")]
692    pub is_virtual: Option<bool>,
693    #[serde(skip_serializing_if = "Option::is_none")]
694    pub is_meta: Option<bool>,
695    #[serde(skip_serializing_if = "Option::is_none")]
696    pub advisor_model: Option<String>,
697    pub uuid: String,
698    pub timestamp: String,
699    #[serde(skip_serializing_if = "Option::is_none")]
700    pub parent_uuid: Option<String>,
701}
702
703/// Derive UUID from parent UUID and index
704pub fn derive_uuid(parent_uuid: &str, index: usize) -> String {
705    let hex = format!("{:012x}", index);
706    let parent_trimmed = parent_uuid.replace('-', "");
707    let prefix = &parent_trimmed[..24.min(parent_trimmed.len())];
708    format!("{}-{}-{}", &prefix[0..8], &prefix[8..12], hex)
709}
710
711/// Get tool use ID from a message
712pub fn get_tool_use_id(message: &NormalizedMessage) -> Option<String> {
713    match message {
714        NormalizedMessage::Assistant(msg) => {
715            if let Some(first) = msg.message.content.first() {
716                if let Ok(block) = serde_json::from_value::<ContentBlock>(first.clone()) {
717                    match block {
718                        ContentBlock::ToolUse { id, .. } => Some(id),
719                        _ => None,
720                    }
721                } else {
722                    // Try to extract id from raw JSON
723                    first.get("id").and_then(|v| v.as_str()).map(String::from)
724                }
725            } else {
726                None
727            }
728        }
729        _ => None,
730    }
731}
732
733/// Message lookups for efficient O(1) access
734#[derive(Debug, Default)]
735pub struct MessageLookups {
736    pub sibling_tool_use_ids: HashMap<String, HashSet<String>>,
737    pub progress_messages_by_tool_use_id: HashMap<String, Vec<ProgressMessage>>,
738    pub in_progress_hook_counts: HashMap<String, HashMap<String, u32>>,
739    pub resolved_hook_counts: HashMap<String, HashMap<String, u32>>,
740    pub tool_result_by_tool_use_id: HashMap<String, NormalizedMessage>,
741    pub tool_use_by_tool_use_id: HashMap<String, serde_json::Value>,
742    pub normalized_message_count: usize,
743    pub resolved_tool_use_ids: HashSet<String>,
744    pub errored_tool_use_ids: HashSet<String>,
745}
746
747/// Build message lookups from normalized messages
748pub fn build_message_lookups(
749    normalized_messages: &[NormalizedMessage],
750    messages: &[Message],
751) -> MessageLookups {
752    let mut lookups = MessageLookups::default();
753
754    // First pass: collect tool use IDs by message ID
755    let mut tool_use_ids_by_message_id: HashMap<String, HashSet<String>> = HashMap::new();
756    let mut tool_use_id_to_message_id: HashMap<String, String> = HashMap::new();
757    let mut tool_use_by_tool_use_id: HashMap<String, serde_json::Value> = HashMap::new();
758
759    for msg in messages {
760        if let Message::Assistant(assistant) = msg {
761            let id = &assistant.message.id;
762            let mut tool_use_ids = HashSet::new();
763            for content in &assistant.message.content {
764                if let Ok(block) = serde_json::from_value::<ContentBlock>(content.clone()) {
765                    if let ContentBlock::ToolUse { id: tool_id, .. } = block {
766                        tool_use_ids.insert(tool_id.clone());
767                        tool_use_id_to_message_id.insert(tool_id.clone(), id.clone());
768                        tool_use_by_tool_use_id.insert(tool_id.clone(), content.clone());
769                    }
770                }
771            }
772            if !tool_use_ids.is_empty() {
773                tool_use_ids_by_message_id.insert(id.clone(), tool_use_ids);
774            }
775        }
776    }
777
778    // Build sibling lookup
779    for (tool_use_id, message_id) in &tool_use_id_to_message_id {
780        if let Some(ids) = tool_use_ids_by_message_id.get(message_id) {
781            lookups
782                .sibling_tool_use_ids
783                .insert(tool_use_id.clone(), ids.clone());
784        }
785    }
786
787    // Second pass: build progress, hook, and tool result lookups
788    for msg in normalized_messages {
789        if let NormalizedMessage::Progress(progress) = msg {
790            let tool_use_id = progress.parent_tool_use_id.clone().unwrap_or_default();
791            if !tool_use_id.is_empty() {
792                lookups
793                    .progress_messages_by_tool_use_id
794                    .entry(tool_use_id.clone())
795                    .or_insert_with(Vec::new)
796                    .push(progress.clone());
797            }
798        }
799
800        // Tool result lookup
801        if let NormalizedMessage::User(user) = msg {
802            if let MessageContent::Blocks(blocks) = &user.message {
803                for block in blocks {
804                    if let ContentBlock::ToolResult {
805                        tool_use_id,
806                        is_error,
807                        ..
808                    } = block
809                    {
810                        lookups.resolved_tool_use_ids.insert(tool_use_id.clone());
811                        if is_error == &Some(true) {
812                            lookups.errored_tool_use_ids.insert(tool_use_id.clone());
813                        }
814                    }
815                }
816            }
817        }
818
819        // Server tool results
820        if let NormalizedMessage::Assistant(assistant) = msg {
821            for content in &assistant.message.content {
822                // Check for server_tool_use, mcp_tool_use
823                if let Some(tool_use_id) = content.get("id") {
824                    if let Some(id_str) = tool_use_id.as_str() {
825                        // Check if there's a corresponding result
826                        let has_result = lookups.resolved_tool_use_ids.contains(id_str);
827                        if !has_result {
828                            // Mark as resolved but not errored
829                            lookups.resolved_tool_use_ids.insert(id_str.to_string());
830                        }
831                    }
832                }
833            }
834        }
835    }
836
837    lookups.tool_use_by_tool_use_id = tool_use_by_tool_use_id;
838    lookups.normalized_message_count = normalized_messages.len();
839
840    lookups
841}
842
843/// Get sibling tool use IDs from lookup
844pub fn get_sibling_tool_use_ids_from_lookup(
845    message: &NormalizedMessage,
846    lookups: &MessageLookups,
847) -> HashSet<String> {
848    let tool_use_id = match get_tool_use_id(message) {
849        Some(id) => id,
850        None => return HashSet::new(),
851    };
852    lookups
853        .sibling_tool_use_ids
854        .get(&tool_use_id)
855        .cloned()
856        .unwrap_or_default()
857}
858
859/// Get progress messages from lookup
860pub fn get_progress_messages_from_lookup(
861    message: &NormalizedMessage,
862    lookups: &MessageLookups,
863) -> Vec<ProgressMessage> {
864    let tool_use_id = match get_tool_use_id(message) {
865        Some(id) => id,
866        None => return Vec::new(),
867    };
868    lookups
869        .progress_messages_by_tool_use_id
870        .get(&tool_use_id)
871        .cloned()
872        .unwrap_or_default()
873}
874
875/// Get tool result IDs from normalized messages
876pub fn get_tool_result_ids(normalized_messages: &[NormalizedMessage]) -> HashMap<String, bool> {
877    let mut result = HashMap::new();
878
879    for msg in normalized_messages {
880        if let NormalizedMessage::User(user) = msg {
881            if let MessageContent::Blocks(blocks) = &user.message {
882                for block in blocks {
883                    if let ContentBlock::ToolResult {
884                        tool_use_id,
885                        is_error,
886                        ..
887                    } = block
888                    {
889                        result.insert(tool_use_id.clone(), is_error.unwrap_or(false));
890                    }
891                }
892            }
893        }
894    }
895
896    result
897}
898
899/// Reorder attachments for API (bubble up until hitting tool result or assistant)
900pub fn reorder_attachments_for_api(messages: Vec<Message>) -> Vec<Message> {
901    let mut result = Vec::new();
902    let mut pending_attachments: Vec<Message> = Vec::new();
903
904    // Scan from bottom to top
905    for i in (0..messages.len()).rev() {
906        let message = messages[i].clone();
907
908        if let Message::Attachment(_) = message {
909            pending_attachments.push(message);
910        } else {
911            let is_stopping_point = matches!(
912                message,
913                Message::Assistant(_) | Message::User(_) if has_tool_result(&message)
914            );
915
916            if is_stopping_point && !pending_attachments.is_empty() {
917                // Reverse pending attachments to maintain order
918                for att in pending_attachments.drain(..).rev() {
919                    result.push(att);
920                }
921                result.push(message);
922            } else {
923                result.push(message);
924            }
925        }
926    }
927
928    // Remaining attachments go to the top
929    for att in pending_attachments.drain(..).rev() {
930        result.push(att);
931    }
932
933    result.reverse();
934    result
935}
936
937/// Check if message has tool result
938fn has_tool_result(message: &Message) -> bool {
939    if let Message::User(user) = message {
940        if let MessageContent::Blocks(blocks) = &user.message {
941            return blocks
942                .iter()
943                .any(|b| matches!(b, ContentBlock::ToolResult { .. }));
944        }
945    }
946    false
947}
948
949/// Check if message is a tool use request
950pub fn is_tool_use_request_message(message: &Message) -> bool {
951    if let Message::Assistant(assistant) = message {
952        assistant.message.content.iter().any(|c| {
953            if let Ok(block) = serde_json::from_value::<ContentBlock>(c.clone()) {
954                matches!(block, ContentBlock::ToolUse { .. })
955            } else {
956                c.get("type").and_then(|t| t.as_str()) == Some("tool_use")
957            }
958        })
959    } else {
960        false
961    }
962}
963
964/// Check if message is a tool result message
965pub fn is_tool_use_result_message(message: &Message) -> bool {
966    if let Message::User(user) = message {
967        if let MessageContent::Blocks(blocks) = &user.message {
968            return blocks
969                .iter()
970                .any(|b| matches!(b, ContentBlock::ToolResult { .. }));
971        }
972    }
973    false
974}
975
976/// Get last assistant message
977pub fn get_last_assistant_message(messages: &[Message]) -> Option<&AssistantMessage> {
978    messages.iter().rev().find_map(|m| {
979        if let Message::Assistant(a) = m {
980            Some(a)
981        } else {
982            None
983        }
984    })
985}
986
987/// Check if last assistant turn has tool calls
988pub fn has_tool_calls_in_last_assistant_turn(messages: &[Message]) -> bool {
989    for msg in messages.iter().rev() {
990        if let Message::Assistant(assistant) = msg {
991            return assistant.message.content.iter().any(|c| {
992                if let Ok(block) = serde_json::from_value::<ContentBlock>(c.clone()) {
993                    matches!(block, ContentBlock::ToolUse { .. })
994                } else {
995                    c.get("type").and_then(|t| t.as_str()) == Some("tool_use")
996                }
997            });
998        }
999    }
1000    false
1001}
1002
1003/// Empty lookups for static rendering
1004pub fn empty_lookups() -> MessageLookups {
1005    MessageLookups::default()
1006}
1007
1008/// Empty string set singleton
1009pub fn empty_string_set() -> HashSet<String> {
1010    HashSet::new()
1011}
1012
1013#[cfg(test)]
1014mod tests {
1015    use super::*;
1016
1017    #[test]
1018    fn test_derive_short_message_id() {
1019        let uuid = "550e8400-e29b-41d4-a716-446655440000";
1020        let short_id = derive_short_message_id(uuid);
1021        assert_eq!(short_id.len(), 6);
1022    }
1023
1024    #[test]
1025    fn test_build_yolo_rejection_message() {
1026        let msg = build_yolo_rejection_message("dangerous command");
1027        assert!(msg.contains("Permission for this action has been denied"));
1028        assert!(msg.contains("dangerous command"));
1029    }
1030
1031    #[test]
1032    fn test_is_classifier_denial() {
1033        assert!(is_classifier_denial(
1034            "Permission for this action has been denied. Reason: testing"
1035        ));
1036        assert!(!is_classifier_denial("Just a regular message"));
1037    }
1038
1039    #[test]
1040    fn test_extract_tag() {
1041        let html = "<test>Hello World</test>";
1042        let extracted = extract_tag(html, "test");
1043        assert_eq!(extracted, Some("Hello World".to_string()));
1044    }
1045
1046    #[test]
1047    fn test_derive_uuid() {
1048        let parent = "550e8400-e29b-41d4-a716-446655440000";
1049        let derived = derive_uuid(parent, 0);
1050        assert!(!derived.is_empty());
1051    }
1052
1053    #[test]
1054    fn test_get_tool_result_ids() {
1055        let messages = vec![NormalizedMessage::User(NormalizedUserMessage {
1056            message: MessageContent::Blocks(vec![ContentBlock::ToolResult {
1057                tool_use_id: "test-id".to_string(),
1058                content: None,
1059                is_error: Some(false),
1060            }]),
1061            extra: UserMessageExtra {
1062                uuid: "uuid".to_string(),
1063                timestamp: "timestamp".to_string(),
1064                ..Default::default()
1065            },
1066        })];
1067
1068        let ids = get_tool_result_ids(&messages);
1069        assert!(ids.contains_key("test-id"));
1070    }
1071}
ai_agent/utils/messages.rs

ai_agent/utils/
messages.rs
