use super::*;
pub fn marker_slug<'a>(agent: &'a AgentKind, custom_name: Option<&'a str>) -> &'a str {
match (agent, custom_name) {
(AgentKind::Custom, Some(name)) if !name.is_empty() => name,
_ => agent.as_str(),
}
}
pub(crate) fn marker_path(agent: &AgentKind, custom_name: Option<&str>) -> PathBuf {
#[cfg(test)]
assert_marker_path_isolated();
aid_dir().join(format!("rate-limit-{}", marker_slug(agent, custom_name)))
}
pub(crate) fn group_marker_path(
agent: &AgentKind,
custom_name: Option<&str>,
group: &str,
) -> PathBuf {
#[cfg(test)]
assert_marker_path_isolated();
let group = group.to_ascii_lowercase();
aid_dir().join(format!(
"rate-limit-{}--{}",
marker_slug(agent, custom_name),
group
))
}
pub fn mark_group_rate_limited(
agent: &AgentKind,
custom_name: Option<&str>,
group: &str,
message: &str,
) {
let provider = (*agent == AgentKind::OpenCode).then_some(group);
write_marker(
&group_marker_path(agent, custom_name, group),
message,
provider,
);
}
pub fn is_group_rate_limited(agent: &AgentKind, custom_name: Option<&str>, group: &str) -> bool {
matches!(
availability_for_group(agent, custom_name, group).status,
RouteStatus::Held
)
}
pub fn clear_group_rate_limit(agent: &AgentKind, custom_name: Option<&str>, group: &str) -> bool {
fs::remove_file(group_marker_path(agent, custom_name, group)).is_ok()
}
pub fn mark_rate_limited(agent: &AgentKind, custom_name: Option<&str>, message: &str) {
write_marker(&marker_path(agent, custom_name), message, None);
}
pub fn mark_rate_limited_for_message(agent: &AgentKind, custom_name: Option<&str>, message: &str) {
mark_rate_limited_for_model(agent, custom_name, None, message);
}
pub fn mark_rate_limited_for_model(
agent: &AgentKind,
custom_name: Option<&str>,
model: Option<&str>,
message: &str,
) {
if crate::provider_capacity::is_capacity(message) {
mark_group_rate_limited(
agent,
custom_name,
&crate::provider_capacity::group(model),
message,
);
return;
}
match quota_group_for_mark(*agent, model, message, None) {
Some(group) => mark_group_rate_limited(agent, custom_name, group, message),
None => mark_rate_limited(agent, custom_name, message),
}
}
pub fn mark_rate_limited_for_value(
agent: &AgentKind,
custom_name: Option<&str>,
evidence: &serde_json::Value,
message: &str,
) {
mark_rate_limited_for_model_value(agent, custom_name, None, evidence, message);
}
pub fn mark_rate_limited_for_model_value(
agent: &AgentKind,
custom_name: Option<&str>,
model: Option<&str>,
evidence: &serde_json::Value,
message: &str,
) {
match quota_group_for_mark(*agent, model, message, Some(evidence)) {
Some(group) => mark_group_rate_limited(agent, custom_name, group, message),
None => mark_rate_limited(agent, custom_name, message),
}
}
fn quota_group_for_mark<'a>(
agent: AgentKind,
model: Option<&'a str>,
message: &'a str,
evidence: Option<&'a serde_json::Value>,
) -> Option<&'a str> {
if agent == AgentKind::Droid {
return crate::agent::model_group::group_from_refusal(agent, message);
}
if let Some(name) = model
&& let Some(group) = crate::agent::model_group::model_group(agent, Some(name))
{
return Some(group);
}
evidence
.and_then(|value| crate::agent::model_group::group_from_refusal_value(agent, value))
.or_else(|| crate::agent::model_group::group_from_refusal(agent, message))
}
pub(crate) fn format_recovery(at: NaiveDateTime) -> String {
at.format("%b %d, %Y %I:%M %p").to_string()
}
fn write_marker(path: &std::path::Path, message: &str, provider: Option<&str>) {
if let Some(parent) = path.parent() {
let _ = fs::create_dir_all(parent);
}
let (recovery_at, hold_line) = match classify_hold(message) {
Hold::Until(at) => (at, String::new()),
Hold::NeedsHuman | Hold::Windowed => (String::new(), format!("hold: {MANUAL_HOLD}\n")),
Hold::Transient => (String::new(), String::new()),
};
let provider_line = format!("provider: {}\n", provider.unwrap_or("unknown"));
let content =
format!("recovery_at: {recovery_at}\n{hold_line}{provider_line}message: {message}\n");
let _ = fs::write(path, content);
}
pub(crate) fn marker_field(content: &str, key: &str) -> Option<String> {
content
.lines()
.find_map(|line| line.strip_prefix(key))
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string)
}