feat(423): auto-schedule timer on rate limit to resume after reset
- pty.rs: detect rate_limit_event hard blocks, parse reset_at, emit WatcherEvent::RateLimitHardBlock with story_id, agent_name, reset_at - watcher.rs: add RateLimitHardBlock variant to WatcherEvent enum - timer.rs: add TimerStore::upsert (add-or-update-to-later) and spawn_rate_limit_auto_scheduler (listens for RateLimitHardBlock, upserts timer for the blocked story) - notifications.rs: handle RateLimitHardBlock events with a debounced chat notification including the scheduled resume time; add format_rate_limit_hard_block_notification helper - matrix/mod.rs: subscribe second watcher_rx for auto-scheduler, pass it to run_bot - matrix/bot/run.rs: wire spawn_rate_limit_auto_scheduler into bot startup Tests cover: AC1 (hard block detection in pty), AC2 (auto-scheduler adds timer), AC3 (upsert deduplication), AC5 (chat notification sent), AC6 (worktree preserved — timer fires start_agent on existing worktree)
This commit is contained in:
+102
-7
@@ -347,13 +347,49 @@ fn run_agent_pty_blocking(
|
||||
// The raw JSON is still forwarded as AgentJson below.
|
||||
"assistant" | "user" => {}
|
||||
"rate_limit_event" => {
|
||||
slog!(
|
||||
"[agent:{story_id}:{agent_name}] API rate limit warning received"
|
||||
);
|
||||
let _ = watcher_tx.send(WatcherEvent::RateLimitWarning {
|
||||
story_id: story_id.to_string(),
|
||||
agent_name: agent_name.to_string(),
|
||||
});
|
||||
let rate_limit_info = json.get("rate_limit_info");
|
||||
let status = rate_limit_info
|
||||
.and_then(|i| i.get("status"))
|
||||
.and_then(|s| s.as_str())
|
||||
.unwrap_or("");
|
||||
let is_hard_block = !status.is_empty() && status != "allowed_warning";
|
||||
let reset_at = rate_limit_info
|
||||
.and_then(|i| i.get("reset_at"))
|
||||
.and_then(|r| r.as_str())
|
||||
.and_then(|r| chrono::DateTime::parse_from_rfc3339(r).ok())
|
||||
.map(|dt| dt.with_timezone(&chrono::Utc));
|
||||
|
||||
if is_hard_block {
|
||||
if let Some(reset_at) = reset_at {
|
||||
slog!(
|
||||
"[agent:{story_id}:{agent_name}] API rate limit hard block \
|
||||
(status={status}); resets at {reset_at}"
|
||||
);
|
||||
let _ = watcher_tx.send(WatcherEvent::RateLimitHardBlock {
|
||||
story_id: story_id.to_string(),
|
||||
agent_name: agent_name.to_string(),
|
||||
reset_at,
|
||||
});
|
||||
} else {
|
||||
slog!(
|
||||
"[agent:{story_id}:{agent_name}] API rate limit hard block \
|
||||
(status={status}); no reset_at in rate_limit_info"
|
||||
);
|
||||
let _ = watcher_tx.send(WatcherEvent::RateLimitWarning {
|
||||
story_id: story_id.to_string(),
|
||||
agent_name: agent_name.to_string(),
|
||||
});
|
||||
}
|
||||
} else {
|
||||
slog!(
|
||||
"[agent:{story_id}:{agent_name}] API rate limit warning received \
|
||||
(status={status})"
|
||||
);
|
||||
let _ = watcher_tx.send(WatcherEvent::RateLimitWarning {
|
||||
story_id: story_id.to_string(),
|
||||
agent_name: agent_name.to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
"result" => {
|
||||
// Extract token usage from the result event.
|
||||
@@ -468,6 +504,65 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
/// AC1: hard block with `reset_at` emits `RateLimitHardBlock` with the
|
||||
/// correct story_id, agent_name, and parsed reset_at timestamp.
|
||||
#[tokio::test]
|
||||
async fn rate_limit_hard_block_sends_watcher_hard_block_event() {
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let script = tmp.path().join("emit_hard_block.sh");
|
||||
std::fs::write(
|
||||
&script,
|
||||
"#!/bin/sh\nprintf '%s\\n' '{\"type\":\"rate_limit_event\",\"rate_limit_info\":{\"status\":\"hard_block\",\"reset_at\":\"2099-01-01T12:00:00Z\"}}'\n",
|
||||
)
|
||||
.unwrap();
|
||||
std::fs::set_permissions(&script, std::fs::Permissions::from_mode(0o755)).unwrap();
|
||||
|
||||
let (tx, _rx) = broadcast::channel::<AgentEvent>(64);
|
||||
let (watcher_tx, mut watcher_rx) = broadcast::channel::<WatcherEvent>(16);
|
||||
let event_log = Arc::new(Mutex::new(Vec::new()));
|
||||
let child_killers = Arc::new(Mutex::new(HashMap::new()));
|
||||
|
||||
let result = run_agent_pty_streaming(
|
||||
"423_story_rate_limit",
|
||||
"coder-1",
|
||||
"sh",
|
||||
&[script.to_string_lossy().to_string()],
|
||||
"--",
|
||||
"/tmp",
|
||||
&tx,
|
||||
&event_log,
|
||||
None,
|
||||
0,
|
||||
child_killers,
|
||||
watcher_tx,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(result.is_ok(), "PTY run should succeed: {:?}", result.err());
|
||||
|
||||
let evt = watcher_rx
|
||||
.try_recv()
|
||||
.expect("Expected a RateLimitHardBlock to be sent on watcher_tx");
|
||||
match evt {
|
||||
WatcherEvent::RateLimitHardBlock {
|
||||
story_id,
|
||||
agent_name,
|
||||
reset_at,
|
||||
} => {
|
||||
assert_eq!(story_id, "423_story_rate_limit");
|
||||
assert_eq!(agent_name, "coder-1");
|
||||
assert_eq!(
|
||||
reset_at.to_rfc3339(),
|
||||
"2099-01-01T12:00:00+00:00",
|
||||
"reset_at should match the parsed timestamp"
|
||||
);
|
||||
}
|
||||
other => panic!("Expected RateLimitHardBlock, got: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_emit_event_writes_to_log_writer() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
|
||||
Reference in New Issue
Block a user