aft/
repeat_breaker.rs

1use std::collections::HashMap;
2use std::hash::{Hash, Hasher};
3use std::time::{Duration, Instant};
4
5use parking_lot::Mutex;
6use serde_json::{Map, Value};
7
8const FIRE_COUNT: u64 = 3;
9const ESCALATE_COUNT: u64 = 6;
10const MIN_SPAN: Duration = Duration::from_secs(30);
11
12#[derive(Debug, Clone, PartialEq, Eq)]
13pub struct RepeatIntervention {
14    pub tool: String,
15    pub count: u64,
16    pub span: Duration,
17}
18
19#[derive(Debug)]
20struct SessionRepeatState {
21    semantic_key: String,
22    output_hash: u64,
23    first_seen: Instant,
24    count: u64,
25}
26
27#[derive(Debug, Default)]
28pub struct RepeatBreaker {
29    sessions: Mutex<HashMap<String, SessionRepeatState>>,
30}
31
32impl RepeatBreaker {
33    pub fn observe(
34        &self,
35        session_id: &str,
36        tool: &str,
37        semantic_key: String,
38        output_hash: u64,
39    ) -> Option<RepeatIntervention> {
40        self.observe_at(session_id, tool, semantic_key, output_hash, Instant::now())
41    }
42
43    #[doc(hidden)]
44    pub fn observe_at(
45        &self,
46        session_id: &str,
47        tool: &str,
48        semantic_key: String,
49        output_hash: u64,
50        now: Instant,
51    ) -> Option<RepeatIntervention> {
52        let mut sessions = self.sessions.lock();
53        let state = sessions
54            .entry(session_id.to_string())
55            .or_insert_with(|| SessionRepeatState {
56                semantic_key: semantic_key.clone(),
57                output_hash,
58                first_seen: now,
59                count: 0,
60            });
61
62        if state.semantic_key != semantic_key || state.output_hash != output_hash {
63            *state = SessionRepeatState {
64                semantic_key,
65                output_hash,
66                first_seen: now,
67                count: 1,
68            };
69            return None;
70        }
71
72        state.count = state.count.saturating_add(1);
73        let span = now.saturating_duration_since(state.first_seen);
74        if state.count < FIRE_COUNT || span < MIN_SPAN {
75            return None;
76        }
77
78        Some(RepeatIntervention {
79            tool: tool.to_string(),
80            count: state.count,
81            span,
82        })
83    }
84
85    pub fn clear_session(&self, session_id: &str) {
86        self.sessions.lock().remove(session_id);
87    }
88
89    pub fn clear(&self) {
90        self.sessions.lock().clear();
91    }
92}
93
94pub fn output_hash(rendered_text: &str) -> u64 {
95    let mut hasher = std::collections::hash_map::DefaultHasher::new();
96    rendered_text.hash(&mut hasher);
97    hasher.finish()
98}
99
100pub fn semantic_key(tool: &str, input: &Value) -> String {
101    let selected = match tool {
102        "bash" | "powershell" => select_fields(input, &["command", "workdir"]),
103        "read" => select_fields(
104            input,
105            &[
106                "path",
107                "filePath",
108                "startLine",
109                "endLine",
110                "offset",
111                "limit",
112            ],
113        ),
114        "grep" => select_fields(
115            input,
116            &[
117                "pattern",
118                "path",
119                "include",
120                "topK",
121                "offset",
122                "includeTests",
123            ],
124        ),
125        "glob" => select_fields(input, &["pattern", "path", "topK", "offset"]),
126        "aft_search" => select_fields(input, &["query", "path", "topK", "offset", "includeTests"]),
127        _ => {
128            let mut value = input.clone();
129            if let Some(object) = value.as_object_mut() {
130                object.remove("description");
131            }
132            value
133        }
134    };
135
136    serde_json::to_string(&(tool, canonicalize(selected))).unwrap_or_default()
137}
138
139fn select_fields(input: &Value, fields: &[&str]) -> Value {
140    let mut selected = Map::new();
141    if let Some(input) = input.as_object() {
142        for field in fields {
143            if let Some(value) = input.get(*field) {
144                selected.insert((*field).to_string(), value.clone());
145            }
146        }
147    }
148    Value::Object(selected)
149}
150
151fn canonicalize(value: Value) -> Value {
152    match value {
153        Value::Object(object) => {
154            let mut entries: Vec<_> = object.into_iter().collect();
155            entries.sort_by(|(left, _), (right, _)| left.cmp(right));
156            Value::Object(
157                entries
158                    .into_iter()
159                    .map(|(key, value)| (key, canonicalize(value)))
160                    .collect(),
161            )
162        }
163        Value::Array(values) => Value::Array(values.into_iter().map(canonicalize).collect()),
164        other => other,
165    }
166}
167
168pub fn escalation_starts_at(count: u64) -> bool {
169    count >= ESCALATE_COUNT
170}
aft/repeat_breaker.rs

aft/
repeat_breaker.rs
