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,
309        reason,
310        DENIAL_WORKAROUND_GUIDANCE
311    )
312}
313
314/// Build message for when classifier is unavailable
315pub fn build_classifier_unavailable_message(tool_name: &str, classifier_model: &str) -> String {
316    format!(
317        "{} 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.",
318        classifier_model, tool_name
319    )
320}
321
322/// Check if tool result message is a classifier denial
323pub fn is_classifier_denial(content: &str) -> bool {
324    content.starts_with(AUTO_MODE_REJECTION_PREFIX)
325}
326
327/// Auto reject message
328pub fn auto_reject_message(tool_name: &str) -> String {
329    format!(
330        "Permission to use {} has been denied. {}",
331        tool_name, DENIAL_WORKAROUND_GUIDANCE
332    )
333}
334
335/// Don't ask reject message
336pub fn dont_ask_reject_message(tool_name: &str) -> String {
337    format!(
338        "Permission to use {} has been denied because Claude Code is running in don't ask mode. {}",
339        tool_name, DENIAL_WORKAROUND_GUIDANCE
340    )
341}
342
343/// Derive short message ID (6-char base36) from UUID
344pub fn derive_short_message_id(uuid: &str) -> String {
345    // Take first 10 hex chars from the UUID (skipping dashes)
346    let hex: String = uuid.replace('-', "").chars().take(10).collect();
347    // Convert to base36 for shorter representation, take 6 chars
348    let parsed = u64::from_str_radix(&hex, 16).unwrap_or(0);
349    let base36 = format!("{:36}", parsed);
350    base36.chars().take(6).collect()
351}
352
353/// Create a user message
354pub fn create_user_message(content: impl Into<MessageContent>) -> UserMessage {
355    let content = content.into();
356    UserMessage {
357        message: content,
358        is_meta: None,
359        is_visible_in_transcript_only: None,
360        is_virtual: None,
361        is_compact_summary: None,
362        summarize_metadata: None,
363        tool_use_result: None,
364        mcp_meta: None,
365        uuid: uuid::Uuid::new_v4().to_string(),
366        timestamp: chrono::Utc::now().to_rfc3339(),
367        image_paste_ids: None,
368        source_tool_assistant_uuid: None,
369        permission_mode: None,
370        origin: None,
371    }
372}
373
374/// Create assistant message
375pub fn create_assistant_message(content: Vec<serde_json::Value>) -> AssistantMessage {
376    AssistantMessage {
377        message: AssistantMessageContent {
378            id: uuid::Uuid::new_v4().to_string(),
379            container: None,
380            model: SYNTHETIC_MODEL.to_string(),
381            role: "assistant".to_string(),
382            stop_reason: Some("stop_sequence".to_string()),
383            stop_sequence: Some("".to_string()),
384            message_type: "message".to_string(),
385            usage: Some(Usage::default()),
386            content,
387            context_management: None,
388        },
389        request_id: None,
390        api_error: None,
391        error: None,
392        error_details: None,
393        is_api_error_message: Some(false),
394        is_virtual: None,
395        is_meta: None,
396        advisor_model: None,
397        uuid: uuid::Uuid::new_v4().to_string(),
398        timestamp: chrono::Utc::now().to_rfc3339(),
399        parent_uuid: None,
400    }
401}
402
403/// Create progress message
404pub fn create_progress_message(
405    tool_use_id: &str,
406    parent_tool_use_id: &str,
407    data: serde_json::Value,
408) -> ProgressMessage {
409    ProgressMessage {
410        data_type: "progress".to_string(),
411        data,
412        tool_use_id: Some(tool_use_id.to_string()),
413        parent_tool_use_id: Some(parent_tool_use_id.to_string()),
414        uuid: uuid::Uuid::new_v4().to_string(),
415        timestamp: chrono::Utc::now().to_rfc3339(),
416        parent_uuid: None,
417    }
418}
419
420/// Create tool result stop message
421pub fn create_tool_result_stop_message(tool_use_id: &str) -> serde_json::Value {
422    serde_json::json!({
423        "type": "tool_result",
424        "content": CANCEL_MESSAGE,
425        "is_error": true,
426        "tool_use_id": tool_use_id
427    })
428}
429
430/// Extract tag from HTML-like content
431pub fn extract_tag(html: &str, tag_name: &str) -> Option<String> {
432    use regex::Regex;
433
434    if html.trim().is_empty() || tag_name.trim().is_empty() {
435        return None;
436    }
437
438    let escaped_tag = tag_name.replace(
439        [
440            '.', '*', '+', '?', '^', '$', '{', '}', '[', ']', '(', ')', '|', '\\',
441        ],
442        "\\$&",
443    );
444
445    let pattern = format!(
446        r"<{}(?:\s+[^>]*)?>([\s\S]*?)</{}>",
447        escaped_tag, escaped_tag
448    );
449
450    let re = Regex::new(&pattern).ok()?;
451
452    let mut depth = 0i32;
453    let mut last_index = 0;
454
455    let opening_tag_re = Regex::new(&format!(r"<{}(?:\s+[^>]*)?>", escaped_tag)).ok()?;
456    let closing_tag_re = Regex::new(&format!(r"</{}>", escaped_tag)).ok()?;
457
458    for caps in re.captures_iter(html) {
459        let content = caps.get(1)?.as_str();
460        let start = caps.get(0)?.start();
461
462        depth = 0;
463
464        for _ in opening_tag_re.find_iter(&html[..start]) {
465            depth += 1;
466        }
467
468        for _ in closing_tag_re.find_iter(&html[..start]) {
469            depth -= 1;
470        }
471
472        if depth == 0 && !content.is_empty() {
473            return Some(content.to_string());
474        }
475
476        last_index = start + caps.get(0)?.len();
477    }
478
479    None
480}
481
482/// Check if message is not empty
483pub fn is_not_empty_message(message: &Message) -> bool {
484    match message {
485        Message::Progress(_) | Message::Attachment(_) | Message::System(_) => true,
486        Message::User(user) => {
487            match &user.message {
488                MessageContent::String(s) => !s.trim().is_empty(),
489                MessageContent::Blocks(blocks) => {
490                    if blocks.is_empty() {
491                        return false;
492                    }
493                    // Skip multi-block messages for now
494                    if blocks.len() > 1 {
495                        return true;
496                    }
497                    // Check first block
498                    match &blocks[0] {
499                        ContentBlock::Text { text } => {
500                            !text.trim().is_empty()
501                                && text != NO_RESPONSE_REQUESTED
502                                && text != INTERRUPT_MESSAGE_FOR_TOOL_USE
503                        }
504                        _ => true,
505                    }
506                }
507            }
508        }
509        Message::Assistant(assistant) => !assistant.message.content.is_empty(),
510    }
511}
512
513/// Normalized message types
514#[derive(Debug, Clone, Serialize, Deserialize)]
515#[serde(tag = "type")]
516pub enum NormalizedMessage {
517    User(NormalizedUserMessage),
518    Assistant(NormalizedAssistantMessage),
519    Progress(ProgressMessage),
520    Attachment(AttachmentMessage),
521    System(SystemMessage),
522}
523
524/// Normalized user message
525#[derive(Debug, Clone, Serialize, Deserialize)]
526pub struct NormalizedUserMessage {
527    pub message: MessageContent,
528    #[serde(flatten)]
529    pub extra: UserMessageExtra,
530}
531
532/// Normalized assistant message
533#[derive(Debug, Clone, Serialize, Deserialize)]
534pub struct NormalizedAssistantMessage {
535    pub message: AssistantMessageContent,
536    #[serde(flatten)]
537    pub extra: AssistantMessageExtra,
538}
539
540/// Extra fields for normalized user message
541#[derive(Debug, Clone, Default, Serialize, Deserialize)]
542pub struct UserMessageExtra {
543    #[serde(skip_serializing_if = "Option::is_none")]
544    pub is_meta: Option<bool>,
545    #[serde(skip_serializing_if = "Option::is_none")]
546    pub is_visible_in_transcript_only: Option<bool>,
547    #[serde(skip_serializing_if = "Option::is_none")]
548    pub is_virtual: Option<bool>,
549    #[serde(skip_serializing_if = "Option::is_none")]
550    pub is_compact_summary: Option<bool>,
551    #[serde(skip_serializing_if = "Option::is_none")]
552    pub summarize_metadata: Option<serde_json::Value>,
553    #[serde(skip_serializing_if = "Option::is_none")]
554    pub tool_use_result: Option<serde_json::Value>,
555    #[serde(skip_serializing_if = "Option::is_none")]
556    pub mcp_meta: Option<serde_json::Value>,
557    pub uuid: String,
558    pub timestamp: String,
559    #[serde(skip_serializing_if = "Option::is_none")]
560    pub image_paste_ids: Option<Vec<u32>>,
561    #[serde(skip_serializing_if = "Option::is_none")]
562    pub source_tool_assistant_uuid: Option<String>,
563    #[serde(skip_serializing_if = "Option::is_none")]
564    pub permission_mode: Option<String>,
565    #[serde(skip_serializing_if = "Option::is_none")]
566    pub origin: Option<MessageOrigin>,
567    #[serde(skip_serializing_if = "Option::is_none")]
568    pub parent_uuid: Option<String>,
569}
570
571/// Extra fields for normalized assistant message
572#[derive(Debug, Clone, Default, Serialize, Deserialize)]
573pub struct AssistantMessageExtra {
574    #[serde(skip_serializing_if = "Option::is_none")]
575    pub request_id: Option<String>,
576    #[serde(skip_serializing_if = "Option::is_none")]
577    pub api_error: Option<serde_json::Value>,
578    #[serde(skip_serializing_if = "Option::is_none")]
579    pub error: Option<serde_json::Value>,
580    #[serde(skip_serializing_if = "Option::is_none")]
581    pub error_details: Option<String>,
582    #[serde(skip_serializing_if = "Option::is_none")]
583    pub is_api_error_message: Option<bool>,
584    #[serde(skip_serializing_if = "Option::is_none")]
585    pub is_virtual: Option<bool>,
586    #[serde(skip_serializing_if = "Option::is_none")]
587    pub is_meta: Option<bool>,
588    #[serde(skip_serializing_if = "Option::is_none")]
589    pub advisor_model: Option<String>,
590    pub uuid: String,
591    pub timestamp: String,
592    #[serde(skip_serializing_if = "Option::is_none")]
593    pub parent_uuid: Option<String>,
594}
595
596/// Derive UUID from parent UUID and index
597pub fn derive_uuid(parent_uuid: &str, index: usize) -> String {
598    let hex = format!("{:012x}", index);
599    let parent_trimmed = parent_uuid.replace('-', "");
600    let prefix = &parent_trimmed[..24.min(parent_trimmed.len())];
601    format!("{}-{}-{}", &prefix[0..8], &prefix[8..12], hex)
602}
603
604/// Get tool use ID from a message
605pub fn get_tool_use_id(message: &NormalizedMessage) -> Option<String> {
606    match message {
607        NormalizedMessage::Assistant(msg) => {
608            if let Some(first) = msg.message.content.first() {
609                if let Ok(block) = serde_json::from_value::<ContentBlock>(first.clone()) {
610                    match block {
611                        ContentBlock::ToolUse { id, .. } => Some(id),
612                        _ => None,
613                    }
614                } else {
615                    // Try to extract id from raw JSON
616                    first.get("id").and_then(|v| v.as_str()).map(String::from)
617                }
618            } else {
619                None
620            }
621        }
622        _ => None,
623    }
624}
625
626/// Message lookups for efficient O(1) access
627#[derive(Debug, Default)]
628pub struct MessageLookups {
629    pub sibling_tool_use_ids: HashMap<String, HashSet<String>>,
630    pub progress_messages_by_tool_use_id: HashMap<String, Vec<ProgressMessage>>,
631    pub in_progress_hook_counts: HashMap<String, HashMap<String, u32>>,
632    pub resolved_hook_counts: HashMap<String, HashMap<String, u32>>,
633    pub tool_result_by_tool_use_id: HashMap<String, NormalizedMessage>,
634    pub tool_use_by_tool_use_id: HashMap<String, serde_json::Value>,
635    pub normalized_message_count: usize,
636    pub resolved_tool_use_ids: HashSet<String>,
637    pub errored_tool_use_ids: HashSet<String>,
638}
639
640/// Build message lookups from normalized messages
641pub fn build_message_lookups(
642    normalized_messages: &[NormalizedMessage],
643    messages: &[Message],
644) -> MessageLookups {
645    let mut lookups = MessageLookups::default();
646
647    // First pass: collect tool use IDs by message ID
648    let mut tool_use_ids_by_message_id: HashMap<String, HashSet<String>> = HashMap::new();
649    let mut tool_use_id_to_message_id: HashMap<String, String> = HashMap::new();
650    let mut tool_use_by_tool_use_id: HashMap<String, serde_json::Value> = HashMap::new();
651
652    for msg in messages {
653        if let Message::Assistant(assistant) = msg {
654            let id = &assistant.message.id;
655            let mut tool_use_ids = HashSet::new();
656            for content in &assistant.message.content {
657                if let Ok(block) = serde_json::from_value::<ContentBlock>(content.clone()) {
658                    if let ContentBlock::ToolUse { id: tool_id, .. } = block {
659                        tool_use_ids.insert(tool_id.clone());
660                        tool_use_id_to_message_id.insert(tool_id.clone(), id.clone());
661                        tool_use_by_tool_use_id.insert(tool_id.clone(), content.clone());
662                    }
663                }
664            }
665            if !tool_use_ids.is_empty() {
666                tool_use_ids_by_message_id.insert(id.clone(), tool_use_ids);
667            }
668        }
669    }
670
671    // Build sibling lookup
672    for (tool_use_id, message_id) in &tool_use_id_to_message_id {
673        if let Some(ids) = tool_use_ids_by_message_id.get(message_id) {
674            lookups
675                .sibling_tool_use_ids
676                .insert(tool_use_id.clone(), ids.clone());
677        }
678    }
679
680    // Second pass: build progress, hook, and tool result lookups
681    for msg in normalized_messages {
682        if let NormalizedMessage::Progress(progress) = msg {
683            let tool_use_id = progress.parent_tool_use_id.clone().unwrap_or_default();
684            if !tool_use_id.is_empty() {
685                lookups
686                    .progress_messages_by_tool_use_id
687                    .entry(tool_use_id.clone())
688                    .or_insert_with(Vec::new)
689                    .push(progress.clone());
690            }
691        }
692
693        // Tool result lookup
694        if let NormalizedMessage::User(user) = msg {
695            if let MessageContent::Blocks(blocks) = &user.message {
696                for block in blocks {
697                    if let ContentBlock::ToolResult {
698                        tool_use_id,
699                        is_error,
700                        ..
701                    } = block
702                    {
703                        lookups.resolved_tool_use_ids.insert(tool_use_id.clone());
704                        if is_error == &Some(true) {
705                            lookups.errored_tool_use_ids.insert(tool_use_id.clone());
706                        }
707                    }
708                }
709            }
710        }
711
712        // Server tool results
713        if let NormalizedMessage::Assistant(assistant) = msg {
714            for content in &assistant.message.content {
715                // Check for server_tool_use, mcp_tool_use
716                if let Some(tool_use_id) = content.get("id") {
717                    if let Some(id_str) = tool_use_id.as_str() {
718                        // Check if there's a corresponding result
719                        let has_result = lookups.resolved_tool_use_ids.contains(id_str);
720                        if !has_result {
721                            // Mark as resolved but not errored
722                            lookups.resolved_tool_use_ids.insert(id_str.to_string());
723                        }
724                    }
725                }
726            }
727        }
728    }
729
730    lookups.tool_use_by_tool_use_id = tool_use_by_tool_use_id;
731    lookups.normalized_message_count = normalized_messages.len();
732
733    lookups
734}
735
736/// Get sibling tool use IDs from lookup
737pub fn get_sibling_tool_use_ids_from_lookup(
738    message: &NormalizedMessage,
739    lookups: &MessageLookups,
740) -> HashSet<String> {
741    let tool_use_id = match get_tool_use_id(message) {
742        Some(id) => id,
743        None => return HashSet::new(),
744    };
745    lookups
746        .sibling_tool_use_ids
747        .get(&tool_use_id)
748        .cloned()
749        .unwrap_or_default()
750}
751
752/// Get progress messages from lookup
753pub fn get_progress_messages_from_lookup(
754    message: &NormalizedMessage,
755    lookups: &MessageLookups,
756) -> Vec<ProgressMessage> {
757    let tool_use_id = match get_tool_use_id(message) {
758        Some(id) => id,
759        None => return Vec::new(),
760    };
761    lookups
762        .progress_messages_by_tool_use_id
763        .get(&tool_use_id)
764        .cloned()
765        .unwrap_or_default()
766}
767
768/// Get tool result IDs from normalized messages
769pub fn get_tool_result_ids(normalized_messages: &[NormalizedMessage]) -> HashMap<String, bool> {
770    let mut result = HashMap::new();
771
772    for msg in normalized_messages {
773        if let NormalizedMessage::User(user) = msg {
774            if let MessageContent::Blocks(blocks) = &user.message {
775                for block in blocks {
776                    if let ContentBlock::ToolResult {
777                        tool_use_id,
778                        is_error,
779                        ..
780                    } = block
781                    {
782                        result.insert(tool_use_id.clone(), is_error.unwrap_or(false));
783                    }
784                }
785            }
786        }
787    }
788
789    result
790}
791
792/// Reorder attachments for API (bubble up until hitting tool result or assistant)
793pub fn reorder_attachments_for_api(messages: Vec<Message>) -> Vec<Message> {
794    let mut result = Vec::new();
795    let mut pending_attachments: Vec<Message> = Vec::new();
796
797    // Scan from bottom to top
798    for i in (0..messages.len()).rev() {
799        let message = messages[i].clone();
800
801        if let Message::Attachment(_) = message {
802            pending_attachments.push(message);
803        } else {
804            let is_stopping_point = matches!(
805                message,
806                Message::Assistant(_) | Message::User(_) if has_tool_result(&message)
807            );
808
809            if is_stopping_point && !pending_attachments.is_empty() {
810                // Reverse pending attachments to maintain order
811                for att in pending_attachments.drain(..).rev() {
812                    result.push(att);
813                }
814                result.push(message);
815            } else {
816                result.push(message);
817            }
818        }
819    }
820
821    // Remaining attachments go to the top
822    for att in pending_attachments.drain(..).rev() {
823        result.push(att);
824    }
825
826    result.reverse();
827    result
828}
829
830/// Check if message has tool result
831fn has_tool_result(message: &Message) -> bool {
832    if let Message::User(user) = message {
833        if let MessageContent::Blocks(blocks) = &user.message {
834            return blocks
835                .iter()
836                .any(|b| matches!(b, ContentBlock::ToolResult { .. }));
837        }
838    }
839    false
840}
841
842/// Check if message is a tool use request
843pub fn is_tool_use_request_message(message: &Message) -> bool {
844    if let Message::Assistant(assistant) = message {
845        assistant.message.content.iter().any(|c| {
846            if let Ok(block) = serde_json::from_value::<ContentBlock>(c.clone()) {
847                matches!(block, ContentBlock::ToolUse { .. })
848            } else {
849                c.get("type").and_then(|t| t.as_str()) == Some("tool_use")
850            }
851        })
852    } else {
853        false
854    }
855}
856
857/// Check if message is a tool result message
858pub fn is_tool_use_result_message(message: &Message) -> bool {
859    if let Message::User(user) = message {
860        if let MessageContent::Blocks(blocks) = &user.message {
861            return blocks
862                .iter()
863                .any(|b| matches!(b, ContentBlock::ToolResult { .. }));
864        }
865    }
866    false
867}
868
869/// Get last assistant message
870pub fn get_last_assistant_message(messages: &[Message]) -> Option<&AssistantMessage> {
871    messages.iter().rev().find_map(|m| {
872        if let Message::Assistant(a) = m {
873            Some(a)
874        } else {
875            None
876        }
877    })
878}
879
880/// Check if last assistant turn has tool calls
881pub fn has_tool_calls_in_last_assistant_turn(messages: &[Message]) -> bool {
882    for msg in messages.iter().rev() {
883        if let Message::Assistant(assistant) = msg {
884            return assistant.message.content.iter().any(|c| {
885                if let Ok(block) = serde_json::from_value::<ContentBlock>(c.clone()) {
886                    matches!(block, ContentBlock::ToolUse { .. })
887                } else {
888                    c.get("type").and_then(|t| t.as_str()) == Some("tool_use")
889                }
890            });
891        }
892    }
893    false
894}
895
896/// Empty lookups for static rendering
897pub fn empty_lookups() -> MessageLookups {
898    MessageLookups::default()
899}
900
901/// Empty string set singleton
902pub fn empty_string_set() -> HashSet<String> {
903    HashSet::new()
904}
905
906#[cfg(test)]
907mod tests {
908    use super::*;
909
910    #[test]
911    fn test_derive_short_message_id() {
912        let uuid = "550e8400-e29b-41d4-a716-446655440000";
913        let short_id = derive_short_message_id(uuid);
914        assert_eq!(short_id.len(), 6);
915    }
916
917    #[test]
918    fn test_build_yolo_rejection_message() {
919        let msg = build_yolo_rejection_message("dangerous command");
920        assert!(msg.contains("Permission for this action has been denied"));
921        assert!(msg.contains("dangerous command"));
922    }
923
924    #[test]
925    fn test_is_classifier_denial() {
926        assert!(is_classifier_denial(
927            "Permission for this action has been denied. Reason: testing"
928        ));
929        assert!(!is_classifier_denial("Just a regular message"));
930    }
931
932    #[test]
933    fn test_extract_tag() {
934        let html = "<test>Hello World</test>";
935        let extracted = extract_tag(html, "test");
936        assert_eq!(extracted, Some("Hello World".to_string()));
937    }
938
939    #[test]
940    fn test_derive_uuid() {
941        let parent = "550e8400-e29b-41d4-a716-446655440000";
942        let derived = derive_uuid(parent, 0);
943        assert!(!derived.is_empty());
944    }
945
946    #[test]
947    fn test_get_tool_result_ids() {
948        let messages = vec![NormalizedMessage::User(NormalizedUserMessage {
949            message: MessageContent::Blocks(vec![ContentBlock::ToolResult {
950                tool_use_id: "test-id".to_string(),
951                content: None,
952                is_error: Some(false),
953            }]),
954            extra: UserMessageExtra {
955                uuid: "uuid".to_string(),
956                timestamp: "timestamp".to_string(),
957                ..Default::default()
958            },
959        })];
960
961        let ids = get_tool_result_ids(&messages);
962        assert!(ids.contains_key("test-id"));
963    }
964}
ai_agent/utils/messages.rs

ai_agent/utils/
messages.rs
