Skip to main content

whycodes_session/
edit_stats.rs

1//! Lines added / removed by this session's own file edits.
2//!
3//! `write`, `edit` and `apply_patch` end a successful result with a
4//! `[+N -M]` trailer (`whycodes_core::edit_trailer`); the agent feeds every
5//! tool result through [`EditStats::record`] and `/info` shows the sum. The
6//! counts are block-level (an `edit` counts the whole replaced text, an
7//! overwrite counts the whole old file), so they are labelled approximate
8//! as soon as one of those two tools contributed.
9
10use serde::{Deserialize, Serialize};
11use whycodes_core::types::{ContentBlock, Message, MessageContent, Role};
12
13/// Cumulative line counts from the session's file-mutating tool results.
14#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
15pub struct EditStats {
16    pub added: u64,
17    pub removed: u64,
18    /// True when a block-level count (`edit`, `write`) is included or a
19    /// counted result had lost its trailer (truncated, spilled, failed).
20    pub approximate: bool,
21}
22
23/// Tools whose results carry a `[+N -M]` trailer.
24const 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    /// Fold one tool result in. Tools without a trailer contract are
32    /// ignored; a counted tool whose result lost its trailer marks the
33    /// total approximate instead of silently undercounting.
34    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    /// Rebuild the counts from a persisted transcript: every assistant
51    /// tool call to a counted tool, paired with its result by call id.
52    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    /// `+12 -3 lines`, with the approximate label when it applies.
80    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        // A counted tool without a trailer (error, truncated) is flagged.
138        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        // c3 never got a result: counted as lost, hence approximate.
169        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}