ai_agent/utils/
analyze_context.rs

1// Source: ~/claudecode/openclaudecode/src/utils/analyzeContext.ts
2use std::collections::HashMap;
3
4use crate::types::{Message, MessageRole};
5use crate::services::token_estimation::rough_token_count_estimation_for_content;
6
7/// Per-category token breakdown
8#[derive(Debug, Clone, Default)]
9pub struct TokenStats {
10    pub tool_requests: HashMap<String, u64>,
11    pub tool_results: HashMap<String, u64>,
12    pub human_messages: u64,
13    pub assistant_messages: u64,
14    pub local_command_outputs: u64,
15    pub other: u64,
16    pub attachments: HashMap<String, u64>,
17    pub duplicate_file_reads: HashMap<String, FileReadStats>,
18    pub total: u64,
19}
20
21#[derive(Debug, Clone, Default)]
22pub struct FileReadStats {
23    pub count: u64,
24    pub tokens: u64,
25}
26
27/// Analyze the token distribution across message categories.
28pub fn analyze_context(messages: &[Message]) -> TokenStats {
29    let mut stats = TokenStats::default();
30    let mut tool_ids_to_names: HashMap<String, String> = HashMap::new();
31    let mut read_tool_id_to_file_path: HashMap<String, String> = HashMap::new();
32    let mut seen_file_paths: HashMap<String, u64> = HashMap::new();
33
34    for msg in messages {
35        match msg.role {
36            MessageRole::User => {
37                let tokens = rough_token_count_estimation_for_content(&msg.content) as u64;
38                stats.human_messages += tokens;
39                stats.total += tokens;
40            }
41            MessageRole::Assistant => {
42                let tokens = rough_token_count_estimation_for_content(&msg.content) as u64;
43                stats.assistant_messages += tokens;
44                stats.total += tokens;
45
46                // Track tool calls for result attribution
47                if let Some(ref tool_calls) = msg.tool_calls {
48                    for call in tool_calls {
49                        tool_ids_to_names.insert(call.id.clone(), call.name.clone());
50
51                        // Track Read tool file paths for duplicate detection
52                        if call.name == "Read" {
53                            if let Some(path) = call.arguments.get("file_path").and_then(|p| p.as_str()) {
54                                read_tool_id_to_file_path
55                                    .insert(call.id.clone(), path.to_string());
56
57                                let tokens = rough_token_count_estimation_for_content(&msg.content) as u64;
58                                let entry = seen_file_paths.entry(path.to_string()).or_insert(0);
59                                *entry += 1;
60                                if *entry > 1 {
61                                    stats.duplicate_file_reads.insert(
62                                        path.to_string(),
63                                        FileReadStats {
64                                            count: *entry,
65                                            tokens,
66                                        },
67                                    );
68                                }
69                            }
70                        }
71                    }
72                }
73            }
74            MessageRole::Tool => {
75                let tokens = rough_token_count_estimation_for_content(&msg.content) as u64;
76                stats.total += tokens;
77
78                if let Some(ref tool_call_id) = msg.tool_call_id {
79                    if let Some(tool_name) = tool_ids_to_names.get(tool_call_id) {
80                        *stats.tool_results.entry(tool_name.clone()).or_insert(0) += tokens;
81
82                        // Track tool request tokens too
83                        *stats.tool_requests.entry(tool_name.clone()).or_insert(0) += 10; // overhead estimate
84                    } else {
85                        stats.local_command_outputs += tokens;
86                    }
87                } else {
88                    stats.other += tokens;
89                }
90            }
91            MessageRole::System => {
92                let tokens = rough_token_count_estimation_for_content(&msg.content) as u64;
93                stats.other += tokens;
94                stats.total += tokens;
95            }
96        }
97
98        // Count attachment tokens
99        if let Some(ref attachments) = msg.attachments {
100            for attachment in attachments {
101                let attachment_name = match attachment {
102                    crate::types::Attachment::File { path } => "File".to_string(),
103                    crate::types::Attachment::AlreadyReadFile { path, .. } => "AlreadyReadFile".to_string(),
104                    crate::types::Attachment::PdfReference { .. } => "PdfReference".to_string(),
105                    crate::types::Attachment::EditedTextFile { .. } => "EditedTextFile".to_string(),
106                    crate::types::Attachment::EditedImageFile { .. } => "EditedImageFile".to_string(),
107                    crate::types::Attachment::Directory { .. } => "Directory".to_string(),
108                    crate::types::Attachment::SelectedLinesInIde { .. } => "SelectedLinesInIde".to_string(),
109                    crate::types::Attachment::MemoryFile { .. } => "MemoryFile".to_string(),
110                    crate::types::Attachment::SkillListing { .. } => "SkillListing".to_string(),
111                    crate::types::Attachment::InvokedSkills { .. } => "InvokedSkills".to_string(),
112                    crate::types::Attachment::TaskStatus { .. } => "TaskStatus".to_string(),
113                    crate::types::Attachment::PlanFileReference { .. } => "PlanFileReference".to_string(),
114                    crate::types::Attachment::McpResources { .. } => "McpResources".to_string(),
115                    crate::types::Attachment::DeferredTools { .. } => "DeferredTools".to_string(),
116                    crate::types::Attachment::AgentListing { .. } => "AgentListing".to_string(),
117                    crate::types::Attachment::Custom { name, .. } => name.clone(),
118                };
119                let attachment_tokens = serde_json::to_string(attachment)
120                    .map(|s| rough_token_count_estimation_for_content(&s) as u64)
121                    .unwrap_or(0);
122                *stats.attachments.entry(attachment_name).or_insert(0) += attachment_tokens;
123            }
124        }
125    }
126
127    stats
128}
129
130/// Convert TokenStats to a flat metrics map for analytics.
131pub fn token_stats_to_metrics(stats: &TokenStats) -> HashMap<String, f64> {
132    let mut metrics = HashMap::new();
133    let total = if stats.total > 0 { stats.total as f64 } else { 1.0 };
134
135    metrics.insert("tool_requests_count".to_string(), stats.tool_requests.len() as f64);
136    metrics.insert("tool_results_count".to_string(), stats.tool_results.len() as f64);
137    metrics.insert(
138        "human_messages_pct".to_string(),
139        stats.human_messages as f64 / total * 100.0,
140    );
141    metrics.insert(
142        "assistant_messages_pct".to_string(),
143        stats.assistant_messages as f64 / total * 100.0,
144    );
145    metrics.insert(
146        "local_command_outputs_pct".to_string(),
147        stats.local_command_outputs as f64 / total * 100.0,
148    );
149    metrics.insert(
150        "other_pct".to_string(),
151        stats.other as f64 / total * 100.0,
152    );
153    metrics.insert("attachments_count".to_string(), stats.attachments.len() as f64);
154    metrics.insert(
155        "duplicate_file_reads_count".to_string(),
156        stats.duplicate_file_reads.len() as f64,
157    );
158    metrics.insert("total_tokens".to_string(), stats.total as f64);
159
160    metrics
161}
162
163/// Analyze context usage and return as JSON (for tool/plugin use).
164pub async fn analyze_context_usage() -> Result<serde_json::Value, String> {
165    Ok(serde_json::json!({
166        "message": "Use analyze_context() directly with messages",
167    }))
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    #[test]
175    fn test_analyze_context_empty() {
176        let stats = analyze_context(&[]);
177        assert_eq!(stats.total, 0);
178    }
179
180    #[test]
181    fn test_analyze_context_user_message() {
182        let messages = vec![Message {
183            role: MessageRole::User,
184            content: "Hello world".to_string(),
185            ..Default::default()
186        }];
187        let stats = analyze_context(&messages);
188        assert!(stats.human_messages > 0);
189        assert_eq!(stats.total, stats.human_messages);
190    }
191
192    #[test]
193    fn test_analyze_context_assistant_message() {
194        let messages = vec![Message {
195            role: MessageRole::Assistant,
196            content: "Here's the answer".to_string(),
197            ..Default::default()
198        }];
199        let stats = analyze_context(&messages);
200        assert!(stats.assistant_messages > 0);
201    }
202
203    #[test]
204    fn test_token_stats_to_metrics() {
205        let mut stats = TokenStats::default();
206        stats.total = 1000;
207        stats.human_messages = 300;
208        let metrics = token_stats_to_metrics(&stats);
209        assert!((metrics.get("human_messages_pct").unwrap() - 30.0) < 0.1);
210    }
211}
ai_agent/utils/analyze_context.rs

ai_agent/utils/
analyze_context.rs
