1use rusqlite::{Connection, OptionalExtension, params};
2
3use super::RunStore;
4use crate::db::StoreError;
5
6#[derive(Debug, Clone, PartialEq, Eq)]
7pub struct CooldownUpdate<'a> {
8 pub target: &'a str,
9 pub until_ms: i64,
10 pub reason: &'a str,
11}
12
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub struct CooldownRecord {
15 pub target: String,
16 pub until_ms: i64,
17 pub reason: String,
18}
19
20pub(crate) mod tx {
21 use rusqlite::{Transaction, params};
22
23 use super::CooldownUpdate;
24
25 pub(crate) fn detach_cooldowns_from_run(
26 transaction: &Transaction<'_>,
27 run_id: &str,
28 ) -> rusqlite::Result<usize> {
29 transaction.execute(
30 "UPDATE cooldowns SET source_run_id = NULL WHERE source_run_id = ?1",
31 params![run_id],
32 )
33 }
34
35 pub(crate) fn delete_budget_reservation_for_run(
36 transaction: &Transaction<'_>,
37 run_id: &str,
38 ) -> rusqlite::Result<usize> {
39 transaction.execute(
40 "DELETE FROM budget_reservations WHERE run_id = ?1",
41 params![run_id],
42 )
43 }
44
45 pub(crate) fn upsert_cooldown(
46 transaction: &Transaction<'_>,
47 run_id: &str,
48 cooldown: &CooldownUpdate<'_>,
49 now_ms: i64,
50 ) -> rusqlite::Result<()> {
51 transaction.execute(
52 "INSERT INTO cooldowns (key, until_ms, reason, source_run_id, updated_at_ms)
53 VALUES ('agent_target:' || ?1, ?2, ?3, ?4, ?5)
54 ON CONFLICT(key) DO UPDATE SET
55 until_ms = MAX(cooldowns.until_ms, excluded.until_ms),
56 reason = CASE WHEN excluded.until_ms >= cooldowns.until_ms
57 THEN excluded.reason ELSE cooldowns.reason END,
58 source_run_id = CASE WHEN excluded.until_ms >= cooldowns.until_ms
59 THEN excluded.source_run_id ELSE cooldowns.source_run_id END,
60 updated_at_ms = excluded.updated_at_ms",
61 params![
62 cooldown.target,
63 cooldown.until_ms,
64 cooldown.reason,
65 run_id,
66 now_ms
67 ],
68 )?;
69 Ok(())
70 }
71}
72
73fn active_cooldown_for_target(
74 connection: &Connection,
75 target: &str,
76 now_ms: i64,
77) -> rusqlite::Result<Option<CooldownRecord>> {
78 connection
79 .query_row(
80 "SELECT ?1, until_ms, reason FROM cooldowns
81 WHERE key = 'agent_target:' || ?1 AND until_ms > ?2",
82 params![target, now_ms],
83 |row| {
84 Ok(CooldownRecord {
85 target: row.get(0)?,
86 until_ms: row.get(1)?,
87 reason: row.get(2)?,
88 })
89 },
90 )
91 .optional()
92}
93
94fn active_cooldowns(connection: &Connection, now_ms: i64) -> rusqlite::Result<Vec<CooldownRecord>> {
95 let mut statement = connection.prepare(
96 "SELECT SUBSTR(key, 14), until_ms, reason FROM cooldowns
97 WHERE key LIKE 'agent_target:%' AND until_ms > ?1
98 ORDER BY key",
99 )?;
100 statement
101 .query_map(params![now_ms], |row| {
102 Ok(CooldownRecord {
103 target: row.get(0)?,
104 until_ms: row.get(1)?,
105 reason: row.get(2)?,
106 })
107 })?
108 .collect::<Result<Vec<_>, _>>()
109}
110
111fn next_active_cooldown(connection: &Connection, now_ms: i64) -> rusqlite::Result<Option<i64>> {
112 connection.query_row(
113 "SELECT MIN(until_ms) FROM cooldowns WHERE until_ms > ?1",
114 params![now_ms],
115 |row| row.get(0),
116 )
117}
118
119impl RunStore {
120 pub(crate) fn active_cooldown_for_target(
121 &self,
122 target: &str,
123 now_ms: i64,
124 ) -> Result<Option<CooldownRecord>, StoreError> {
125 active_cooldown_for_target(&self.db.lock(), target, now_ms).map_err(StoreError::from)
126 }
127
128 pub(crate) fn active_cooldowns(&self, now_ms: i64) -> Result<Vec<CooldownRecord>, StoreError> {
129 active_cooldowns(&self.db.lock(), now_ms).map_err(StoreError::from)
130 }
131
132 pub(crate) fn next_active_cooldown(&self, now_ms: i64) -> Result<Option<i64>, StoreError> {
133 next_active_cooldown(&self.db.lock(), now_ms).map_err(StoreError::from)
134 }
135}
136
137#[cfg(test)]
138mod tests {
139 use tempfile::tempdir;
140
141 use super::{CooldownUpdate, tx};
142 use crate::outcome::Outcome;
143 use crate::run_store::test_support::{claim_run, open_seeded, settle_run};
144
145 #[test]
146 fn cooldowns_extend_without_shortening() {
147 let directory = tempdir().unwrap();
148 let store = open_seeded(&directory.path().join("sloop.db"));
149
150 claim_run(&store, "R1", "{}", "{}", 2_000);
151 settle_run(
152 &store,
153 "R1",
154 Some(1),
155 Outcome::RateLimited,
156 &[],
157 Some(&CooldownUpdate {
158 target: "claude",
159 until_ms: 10_000,
160 reason: "rate limit",
161 }),
162 2_100,
163 );
164
165 assert_eq!(store.next_active_cooldown(5_000).unwrap(), Some(10_000));
166 assert_eq!(store.active_cooldowns(5_000).unwrap()[0].target, "claude");
167
168 claim_run(&store, "R2", "{}", "{}", 5_000);
169 settle_run(
170 &store,
171 "R2",
172 Some(1),
173 Outcome::RateLimited,
174 &[],
175 Some(&CooldownUpdate {
176 target: "claude",
177 until_ms: 8_000,
178 reason: "shorter retry",
179 }),
180 5_100,
181 );
182
183 let cooldown = store
184 .active_cooldown_for_target("claude", 6_000)
185 .unwrap()
186 .unwrap();
187 assert_eq!(cooldown.until_ms, 10_000);
188 assert_eq!(cooldown.reason, "rate limit");
189 assert!(store.active_cooldowns(10_000).unwrap().is_empty());
190 }
191
192 #[test]
193 fn run_deletion_detaches_cooldowns_and_deletes_budget_reservations() {
194 let directory = tempdir().unwrap();
195 let store = open_seeded(&directory.path().join("sloop.db"));
196 claim_run(&store, "R1", "{}", "{}", 2_000);
197 settle_run(
198 &store,
199 "R1",
200 Some(1),
201 Outcome::Failed,
202 &[],
203 Some(&CooldownUpdate {
204 target: "claude",
205 until_ms: 10_000,
206 reason: "rate limit",
207 }),
208 2_100,
209 );
210 store
211 .db()
212 .lock()
213 .execute(
214 "INSERT INTO budget_reservations
215 (run_id, reserved_tokens, actual_tokens, state, created_at_ms,
216 reconciled_at_ms)
217 VALUES ('R1', 1000, NULL, 'reserved', 2000, NULL)",
218 [],
219 )
220 .unwrap();
221
222 let database = store.db();
223 let mut connection = database.lock();
224 let transaction = connection.transaction().unwrap();
225 tx::detach_cooldowns_from_run(&transaction, "R1").unwrap();
226 tx::delete_budget_reservation_for_run(&transaction, "R1").unwrap();
227 transaction.commit().unwrap();
228 drop(connection);
229
230 let database = store.db();
231 let connection = database.lock();
232 let source_run_id: Option<String> = connection
233 .query_row(
234 "SELECT source_run_id FROM cooldowns WHERE key = 'agent_target:claude'",
235 [],
236 |row| row.get(0),
237 )
238 .unwrap();
239 let reservations: i64 = connection
240 .query_row("SELECT COUNT(*) FROM budget_reservations", [], |row| {
241 row.get(0)
242 })
243 .unwrap();
244 assert_eq!(source_run_id, None);
245 assert_eq!(reservations, 0);
246 }
247}