awaken_runtime/agent/state/
context_throttle.rs

1use crate::state::StateKey;
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4
5/// Per-key throttle entry: tracks when a context message was last injected.
6#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
7pub struct ThrottleEntry {
8    /// Step number when this key was last injected.
9    pub last_step: usize,
10    /// Hash of the content at last injection (re-inject if content changes).
11    pub content_hash: u64,
12}
13
14/// Throttle state for context message injection.
15#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
16pub struct ContextThrottleMap {
17    pub entries: HashMap<String, ThrottleEntry>,
18}
19
20/// Update for the context throttle state.
21pub enum ContextThrottleUpdate {
22    /// Record that a key was injected at a given step with a content hash.
23    Injected {
24        key: String,
25        step: usize,
26        content_hash: u64,
27    },
28}
29
30/// State key for context message throttle tracking.
31///
32/// Tracks per-key injection history so the loop runner can enforce cooldown rules.
33pub struct ContextThrottleState;
34
35impl StateKey for ContextThrottleState {
36    const KEY: &'static str = "__runtime.context_throttle";
37
38    type Value = ContextThrottleMap;
39    type Update = ContextThrottleUpdate;
40
41    fn apply(value: &mut Self::Value, update: Self::Update) {
42        match update {
43            ContextThrottleUpdate::Injected {
44                key,
45                step,
46                content_hash,
47            } => {
48                value.entries.insert(
49                    key,
50                    ThrottleEntry {
51                        last_step: step,
52                        content_hash,
53                    },
54                );
55            }
56        }
57    }
58}
59
60#[cfg(test)]
61mod tests {
62    use super::*;
63
64    #[test]
65    fn throttle_map_default_is_empty() {
66        let map = ContextThrottleMap::default();
67        assert!(map.entries.is_empty());
68    }
69
70    #[test]
71    fn injected_creates_entry() {
72        let mut map = ContextThrottleMap::default();
73        ContextThrottleState::apply(
74            &mut map,
75            ContextThrottleUpdate::Injected {
76                key: "reminder".into(),
77                step: 5,
78                content_hash: 12345,
79            },
80        );
81        assert_eq!(map.entries.len(), 1);
82        let entry = &map.entries["reminder"];
83        assert_eq!(entry.last_step, 5);
84        assert_eq!(entry.content_hash, 12345);
85    }
86
87    #[test]
88    fn injected_updates_existing_entry() {
89        let mut map = ContextThrottleMap::default();
90        ContextThrottleState::apply(
91            &mut map,
92            ContextThrottleUpdate::Injected {
93                key: "reminder".into(),
94                step: 1,
95                content_hash: 111,
96            },
97        );
98        ContextThrottleState::apply(
99            &mut map,
100            ContextThrottleUpdate::Injected {
101                key: "reminder".into(),
102                step: 5,
103                content_hash: 222,
104            },
105        );
106        assert_eq!(map.entries.len(), 1);
107        let entry = &map.entries["reminder"];
108        assert_eq!(entry.last_step, 5);
109        assert_eq!(entry.content_hash, 222);
110    }
111
112    #[test]
113    fn injected_multiple_keys_independent() {
114        let mut map = ContextThrottleMap::default();
115        ContextThrottleState::apply(
116            &mut map,
117            ContextThrottleUpdate::Injected {
118                key: "a".into(),
119                step: 1,
120                content_hash: 100,
121            },
122        );
123        ContextThrottleState::apply(
124            &mut map,
125            ContextThrottleUpdate::Injected {
126                key: "b".into(),
127                step: 2,
128                content_hash: 200,
129            },
130        );
131        assert_eq!(map.entries.len(), 2);
132        assert_eq!(map.entries["a"].last_step, 1);
133        assert_eq!(map.entries["b"].last_step, 2);
134    }
135
136    #[test]
137    fn throttle_entry_serde_roundtrip() {
138        let entry = ThrottleEntry {
139            last_step: 42,
140            content_hash: 987654321,
141        };
142        let json = serde_json::to_string(&entry).unwrap();
143        let parsed: ThrottleEntry = serde_json::from_str(&json).unwrap();
144        assert_eq!(parsed, entry);
145    }
146
147    #[test]
148    fn throttle_map_serde_roundtrip() {
149        let mut map = ContextThrottleMap::default();
150        ContextThrottleState::apply(
151            &mut map,
152            ContextThrottleUpdate::Injected {
153                key: "k1".into(),
154                step: 3,
155                content_hash: 456,
156            },
157        );
158        let json = serde_json::to_string(&map).unwrap();
159        let parsed: ContextThrottleMap = serde_json::from_str(&json).unwrap();
160        assert_eq!(parsed, map);
161    }
162
163    #[test]
164    fn content_hash_change_detection() {
165        let mut map = ContextThrottleMap::default();
166
167        // First injection
168        ContextThrottleState::apply(
169            &mut map,
170            ContextThrottleUpdate::Injected {
171                key: "reminder".into(),
172                step: 1,
173                content_hash: 111,
174            },
175        );
176
177        // Check if content changed (simulating throttle logic)
178        let entry = &map.entries["reminder"];
179        let new_hash: u64 = 222;
180        let content_changed = entry.content_hash != new_hash;
181        assert!(content_changed, "different hash should indicate change");
182
183        // Same hash
184        let same_hash: u64 = 111;
185        let content_same = entry.content_hash == same_hash;
186        assert!(content_same, "same hash should indicate no change");
187    }
188
189    #[test]
190    fn cooldown_check_logic() {
191        let mut map = ContextThrottleMap::default();
192
193        // Inject at step 1 with cooldown of 3
194        ContextThrottleState::apply(
195            &mut map,
196            ContextThrottleUpdate::Injected {
197                key: "reminder".into(),
198                step: 1,
199                content_hash: 100,
200            },
201        );
202
203        let entry = &map.entries["reminder"];
204        let cooldown = 3usize;
205
206        // Step 2: within cooldown (2 - 1 = 1 < 3)
207        assert!(2usize.saturating_sub(entry.last_step) < cooldown);
208
209        // Step 3: within cooldown (3 - 1 = 2 < 3)
210        assert!(3usize.saturating_sub(entry.last_step) < cooldown);
211
212        // Step 4: cooldown expired (4 - 1 = 3 >= 3)
213        assert!(4usize.saturating_sub(entry.last_step) >= cooldown);
214    }
215}
awaken_runtime/agent/state/context_throttle.rs

awaken_runtime/agent/state/
context_throttle.rs
