whycodes_session/
edit_stats.rs1use serde::{Deserialize, Serialize};
11use whycodes_core::types::{ContentBlock, Message, MessageContent, Role};
12
13#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
15pub struct EditStats {
16 pub added: u64,
17 pub removed: u64,
18 pub approximate: bool,
21}
22
23const COUNTED_TOOLS: [&str; 3] = ["write", "edit", "apply_patch"];
25
26impl EditStats {
27 pub fn is_empty(&self) -> bool {
28 self.added == 0 && self.removed == 0
29 }
30
31 pub fn record(&mut self, tool: &str, content: &str) {
35 if !COUNTED_TOOLS.contains(&tool) {
36 return;
37 }
38 match whycodes_core::parse_edit_trailers(content) {
39 Some((added, removed)) => {
40 self.added += added;
41 self.removed += removed;
42 if tool != "apply_patch" {
43 self.approximate = true;
44 }
45 }
46 None => self.approximate = true,
47 }
48 }
49
50 pub fn from_messages(messages: &[Message]) -> Self {
53 let mut stats = Self::default();
54 for (i, msg) in messages.iter().enumerate() {
55 if msg.role != Role::Assistant {
56 continue;
57 }
58 let MessageContent::Blocks(blocks) = &msg.content else {
59 continue;
60 };
61 for block in blocks {
62 let ContentBlock::ToolUse { id, name, .. } = block else {
63 continue;
64 };
65 if !COUNTED_TOOLS.contains(&name.as_str()) {
66 continue;
67 }
68 let result = messages[i + 1..]
69 .iter()
70 .find(|m| m.role == Role::Tool && m.tool_call_id.as_deref() == Some(id))
71 .and_then(|m| m.content.as_text())
72 .unwrap_or("");
73 stats.record(name, result);
74 }
75 }
76 stats
77 }
78
79 pub fn label(&self) -> String {
81 let mut out = format!("+{} -{} lines", self.added, self.removed);
82 if self.approximate {
83 out.push_str(" (approximate)");
84 }
85 out
86 }
87}
88
89#[cfg(test)]
90mod tests {
91 use super::*;
92
93 fn assistant_call(id: &str, name: &str) -> Message {
94 Message {
95 role: Role::Assistant,
96 content: MessageContent::Blocks(vec![ContentBlock::ToolUse {
97 id: id.into(),
98 name: name.into(),
99 input: serde_json::json!({}),
100 }]),
101 tool_call_id: None,
102 name: None,
103 created_at: None,
104 }
105 }
106
107 fn tool_result(id: &str, text: &str) -> Message {
108 Message {
109 role: Role::Tool,
110 content: MessageContent::Text(text.into()),
111 tool_call_id: Some(id.into()),
112 name: None,
113 created_at: None,
114 }
115 }
116
117 #[test]
118 fn record_sums_counted_tools_only() {
119 let mut s = EditStats::default();
120 assert!(s.is_empty());
121 s.record("bash", "[+99 -99]");
122 assert!(s.is_empty());
123 s.record("apply_patch", "Patch applied to `a` (1 hunk).\n[+2 -1]");
124 assert_eq!(
125 s,
126 EditStats {
127 added: 2,
128 removed: 1,
129 approximate: false
130 }
131 );
132 assert_eq!(s.label(), "+2 -1 lines");
133 s.record("write", "Wrote a.rs ยท 3 lines\n\n[+3 -0]");
134 assert!(s.approximate);
135 assert_eq!(s.label(), "+5 -1 lines (approximate)");
136
137 let mut lost = EditStats::default();
139 lost.record("edit", "Error writing file: denied");
140 assert!(lost.is_empty());
141 assert!(lost.approximate);
142 }
143
144 #[test]
145 fn from_messages_pairs_calls_with_results_by_id() {
146 let messages = vec![
147 Message {
148 role: Role::User,
149 content: MessageContent::Text("hi".into()),
150 tool_call_id: None,
151 name: None,
152 created_at: None,
153 },
154 Message {
155 role: Role::Assistant,
156 content: MessageContent::Text("plain".into()),
157 tool_call_id: None,
158 name: None,
159 created_at: None,
160 },
161 assistant_call("c1", "apply_patch"),
162 assistant_call("c2", "read"),
163 tool_result("c2", "contents"),
164 tool_result("c1", "Patch applied to `a` (1 hunk).\n[+4 -2]"),
165 assistant_call("c3", "edit"),
166 ];
167 let stats = EditStats::from_messages(&messages);
168 assert_eq!(
170 stats,
171 EditStats {
172 added: 4,
173 removed: 2,
174 approximate: true
175 }
176 );
177 assert_eq!(EditStats::from_messages(&[]), EditStats::default());
178 }
179}