1use anyhow::{Context, Result};
2use chrono::{TimeZone, Utc};
3use serde_json::{Map, Value};
4use std::path::Path;
5
6use crate::fs;
7use crate::json;
8use crate::rate_limits::render;
9
10pub fn write_weekly(target_file: &Path, usage_json: &Value) -> Result<()> {
11 if !target_file.is_file() {
12 anyhow::bail!("target file not found");
13 }
14
15 let usage = match render::parse_usage(usage_json) {
16 Some(value) => value,
17 None => return Ok(()),
18 };
19 let values = render::render_values(&usage);
20
21 let (weekly_reset_epoch, non_weekly_reset_epoch) = if values.primary_label == "Weekly" {
22 (
23 values.primary_reset_epoch,
24 Some(values.secondary_reset_epoch),
25 )
26 } else {
27 (
28 values.secondary_reset_epoch,
29 Some(values.primary_reset_epoch),
30 )
31 };
32
33 if weekly_reset_epoch <= 0 {
34 return Ok(());
35 }
36
37 let weekly_reset_iso = epoch_to_iso(weekly_reset_epoch)?;
38 let non_weekly_reset_epoch = non_weekly_reset_epoch.filter(|epoch| *epoch > 0);
39 let non_weekly_reset_iso = non_weekly_reset_epoch.and_then(|epoch| epoch_to_iso(epoch).ok());
40
41 let fetched_at_iso = Utc::now().format("%Y-%m-%dT%H:%M:%SZ").to_string();
42
43 let mut root = json::read_json(target_file).unwrap_or_else(|_| Value::Object(Map::new()));
44 let root_obj = root
45 .as_object_mut()
46 .ok_or_else(|| anyhow::anyhow!("root not object"))?;
47
48 let mut codex_rate_limits = root_obj
49 .get("codex_rate_limits")
50 .and_then(|value| value.as_object())
51 .cloned()
52 .unwrap_or_else(Map::new);
53
54 codex_rate_limits.insert(
55 "weekly_reset_at".to_string(),
56 Value::String(weekly_reset_iso.clone()),
57 );
58 codex_rate_limits.insert(
59 "weekly_reset_at_epoch".to_string(),
60 Value::Number(weekly_reset_epoch.into()),
61 );
62 codex_rate_limits.insert(
63 "weekly_fetched_at".to_string(),
64 Value::String(fetched_at_iso),
65 );
66
67 match (non_weekly_reset_epoch, non_weekly_reset_iso) {
68 (Some(epoch), Some(iso)) => {
69 codex_rate_limits.insert("non_weekly_reset_at".to_string(), Value::String(iso));
70 codex_rate_limits.insert(
71 "non_weekly_reset_at_epoch".to_string(),
72 Value::Number(epoch.into()),
73 );
74 }
75 _ => {
76 codex_rate_limits.remove("non_weekly_reset_at");
77 codex_rate_limits.remove("non_weekly_reset_at_epoch");
78 }
79 }
80
81 root_obj.insert(
82 "codex_rate_limits".to_string(),
83 Value::Object(codex_rate_limits),
84 );
85
86 let out = serde_json::to_vec(&root).context("serialize writeback")?;
87 fs::write_atomic(target_file, &out, fs::SECRET_FILE_MODE)?;
88
89 Ok(())
90}
91
92fn epoch_to_iso(epoch: i64) -> Result<String> {
93 if epoch <= 0 {
94 anyhow::bail!("invalid epoch");
95 }
96 Ok(Utc
97 .timestamp_opt(epoch, 0)
98 .single()
99 .context("epoch")?
100 .format("%Y-%m-%dT%H:%M:%SZ")
101 .to_string())
102}
103
104#[cfg(test)]
105mod tests {
106 use super::{epoch_to_iso, write_weekly};
107 use serde_json::{Value, json};
108 use std::fs;
109 use std::path::Path;
110
111 fn write_json(path: &Path, value: &Value) {
112 let bytes = serde_json::to_vec(value).expect("serialize");
113 fs::write(path, bytes).expect("write json");
114 }
115
116 fn read_json(path: &Path) -> Value {
117 let bytes = fs::read(path).expect("read json");
118 serde_json::from_slice(&bytes).expect("parse json")
119 }
120
121 fn usage_with_weekly_secondary() -> Value {
122 json!({
123 "rate_limit": {
124 "primary_window": {
125 "limit_window_seconds": 18000,
126 "used_percent": 6.0,
127 "reset_at": 1700003600
128 },
129 "secondary_window": {
130 "limit_window_seconds": 604800,
131 "used_percent": 12.0,
132 "reset_at": 1700600000
133 }
134 }
135 })
136 }
137
138 #[test]
139 fn write_weekly_uses_primary_window_when_primary_is_weekly() {
140 let dir = tempfile::TempDir::new().expect("tempdir");
141 let target = dir.path().join("alpha.json");
142 write_json(&target, &json!({ "tokens": { "access_token": "tok" } }));
143
144 let usage = json!({
145 "rate_limit": {
146 "primary_window": {
147 "limit_window_seconds": 604800,
148 "used_percent": 12.0,
149 "reset_at": 1700700000
150 },
151 "secondary_window": {
152 "limit_window_seconds": 18000,
153 "used_percent": 6.0,
154 "reset_at": 1700003600
155 }
156 }
157 });
158
159 write_weekly(&target, &usage).expect("write weekly");
160 let written = read_json(&target);
161 let limits = &written["codex_rate_limits"];
162
163 assert_eq!(limits["weekly_reset_at_epoch"].as_i64(), Some(1700700000));
164 assert_eq!(
165 limits["weekly_reset_at"].as_str(),
166 Some(epoch_to_iso(1700700000).expect("weekly iso").as_str())
167 );
168 assert_eq!(
169 limits["non_weekly_reset_at_epoch"].as_i64(),
170 Some(1700003600)
171 );
172 assert_eq!(
173 limits["non_weekly_reset_at"].as_str(),
174 Some(epoch_to_iso(1700003600).expect("non-weekly iso").as_str())
175 );
176 assert!(limits["weekly_fetched_at"].as_str().is_some());
177 }
178
179 #[test]
180 fn write_weekly_preserves_existing_codex_rate_limits_fields() {
181 let dir = tempfile::TempDir::new().expect("tempdir");
182 let target = dir.path().join("alpha.json");
183 write_json(
184 &target,
185 &json!({
186 "tokens": { "access_token": "tok" },
187 "codex_rate_limits": {
188 "source": "legacy-metadata",
189 "weekly_reset_at_epoch": 111
190 }
191 }),
192 );
193
194 let usage = json!({
195 "rate_limit": {
196 "primary_window": {
197 "limit_window_seconds": 18000,
198 "used_percent": 6.0,
199 "reset_at": 1700003600
200 },
201 "secondary_window": {
202 "limit_window_seconds": 604800,
203 "used_percent": 12.0,
204 "reset_at": 1700600000
205 }
206 }
207 });
208
209 write_weekly(&target, &usage).expect("write weekly");
210 let written = read_json(&target);
211 let limits = &written["codex_rate_limits"];
212
213 assert_eq!(limits["source"].as_str(), Some("legacy-metadata"));
214 assert_eq!(limits["weekly_reset_at_epoch"].as_i64(), Some(1700600000));
215 assert_eq!(
216 limits["non_weekly_reset_at_epoch"].as_i64(),
217 Some(1700003600)
218 );
219 }
220
221 #[test]
222 fn write_weekly_skips_write_when_weekly_epoch_is_non_positive() {
223 let dir = tempfile::TempDir::new().expect("tempdir");
224 let target = dir.path().join("alpha.json");
225 write_json(
226 &target,
227 &json!({
228 "tokens": { "access_token": "tok" },
229 "codex_rate_limits": {
230 "weekly_reset_at_epoch": 111,
231 "weekly_reset_at": "legacy"
232 }
233 }),
234 );
235 let before = read_json(&target);
236
237 let usage = json!({
238 "rate_limit": {
239 "primary_window": {
240 "limit_window_seconds": 18000,
241 "used_percent": 6.0,
242 "reset_at": 1700003600
243 },
244 "secondary_window": {
245 "limit_window_seconds": 604800,
246 "used_percent": 12.0,
247 "reset_at": 0
248 }
249 }
250 });
251
252 write_weekly(&target, &usage).expect("write weekly");
253 let after = read_json(&target);
254 assert_eq!(after, before);
255 }
256
257 #[test]
258 fn write_weekly_clears_non_weekly_fields_when_non_weekly_epoch_is_non_positive() {
259 let dir = tempfile::TempDir::new().expect("tempdir");
260 let target = dir.path().join("alpha.json");
261 write_json(
262 &target,
263 &json!({
264 "tokens": { "access_token": "tok" },
265 "codex_rate_limits": {
266 "source": "legacy-metadata",
267 "non_weekly_reset_at_epoch": 1700003600,
268 "non_weekly_reset_at": "2023-11-14T23:13:20Z"
269 }
270 }),
271 );
272
273 let usage = json!({
274 "rate_limit": {
275 "primary_window": {
276 "limit_window_seconds": 18000,
277 "used_percent": 6.0,
278 "reset_at": 0
279 },
280 "secondary_window": {
281 "limit_window_seconds": 604800,
282 "used_percent": 12.0,
283 "reset_at": 1700600000
284 }
285 }
286 });
287
288 write_weekly(&target, &usage).expect("write weekly");
289 let written = read_json(&target);
290 let limits = written["codex_rate_limits"]
291 .as_object()
292 .expect("limits object");
293
294 assert_eq!(
295 limits.get("source").and_then(Value::as_str),
296 Some("legacy-metadata")
297 );
298 assert_eq!(
299 limits.get("weekly_reset_at_epoch").and_then(Value::as_i64),
300 Some(1700600000)
301 );
302 assert!(
303 limits
304 .get("weekly_reset_at")
305 .and_then(Value::as_str)
306 .is_some()
307 );
308 assert!(
309 limits
310 .get("weekly_fetched_at")
311 .and_then(Value::as_str)
312 .is_some()
313 );
314 assert!(!limits.contains_key("non_weekly_reset_at"));
315 assert!(!limits.contains_key("non_weekly_reset_at_epoch"));
316 }
317
318 #[test]
319 fn write_weekly_fails_when_target_file_is_missing() {
320 let dir = tempfile::TempDir::new().expect("tempdir");
321 let target = dir.path().join("missing.json");
322 let usage = json!({
323 "rate_limit": {
324 "primary_window": {
325 "limit_window_seconds": 18000,
326 "used_percent": 6.0,
327 "reset_at": 1700003600
328 },
329 "secondary_window": {
330 "limit_window_seconds": 604800,
331 "used_percent": 12.0,
332 "reset_at": 1700600000
333 }
334 }
335 });
336
337 let err = write_weekly(&target, &usage).expect_err("missing target must fail");
338 assert!(err.to_string().contains("target file not found"));
339 }
340
341 #[test]
342 fn write_weekly_noops_when_usage_payload_is_unparseable() {
343 let dir = tempfile::TempDir::new().expect("tempdir");
344 let target = dir.path().join("alpha.json");
345 write_json(
346 &target,
347 &json!({
348 "tokens": { "access_token": "tok" },
349 "codex_rate_limits": {
350 "weekly_reset_at_epoch": 111,
351 "weekly_reset_at": "legacy"
352 }
353 }),
354 );
355 let before = read_json(&target);
356
357 write_weekly(&target, &json!({ "unexpected": "shape" })).expect("write weekly");
358 let after = read_json(&target);
359
360 assert_eq!(after, before);
361 }
362
363 #[test]
364 fn write_weekly_recovers_from_malformed_existing_json() {
365 let dir = tempfile::TempDir::new().expect("tempdir");
366 let target = dir.path().join("alpha.json");
367 fs::write(&target, b"{ malformed").expect("write malformed json");
368
369 write_weekly(&target, &usage_with_weekly_secondary()).expect("write weekly");
370 let written = read_json(&target);
371 let limits = written["codex_rate_limits"]
372 .as_object()
373 .expect("limits object");
374
375 assert_eq!(
376 limits.get("weekly_reset_at_epoch").and_then(Value::as_i64),
377 Some(1700600000)
378 );
379 assert_eq!(
380 limits
381 .get("non_weekly_reset_at_epoch")
382 .and_then(Value::as_i64),
383 Some(1700003600)
384 );
385 assert!(
386 limits
387 .get("weekly_fetched_at")
388 .and_then(Value::as_str)
389 .is_some()
390 );
391 }
392
393 #[test]
394 fn write_weekly_fails_when_existing_json_root_is_not_object() {
395 let dir = tempfile::TempDir::new().expect("tempdir");
396 let target = dir.path().join("alpha.json");
397 write_json(&target, &json!(["not", "an", "object"]));
398
399 let err = write_weekly(&target, &usage_with_weekly_secondary())
400 .expect_err("non-object root should fail");
401
402 assert!(err.to_string().contains("root not object"));
403 }
404
405 #[test]
406 fn write_weekly_replaces_non_object_codex_rate_limits_value() {
407 let dir = tempfile::TempDir::new().expect("tempdir");
408 let target = dir.path().join("alpha.json");
409 write_json(
410 &target,
411 &json!({
412 "tokens": { "access_token": "tok" },
413 "codex_rate_limits": "legacy-string"
414 }),
415 );
416
417 write_weekly(&target, &usage_with_weekly_secondary()).expect("write weekly");
418 let written = read_json(&target);
419
420 assert_eq!(written["tokens"]["access_token"].as_str(), Some("tok"));
421 assert!(written["codex_rate_limits"].is_object());
422 assert_eq!(
423 written["codex_rate_limits"]["weekly_reset_at_epoch"].as_i64(),
424 Some(1700600000)
425 );
426 }
427}