tuitbot_core/storage/rate_limits/
mod.rs

1//! Rate limit tracking and enforcement for action quotas.
2//!
3//! Tracks usage of actions (replies, tweets, threads) within time periods,
4//! stored in SQLite for persistence across restarts. Check and reset logic
5//! uses transactions for atomicity.
6
7pub mod queries;
8pub mod tracker;
9
10pub use crate::mcp_policy::types::{PolicyRateLimit, RateLimitDimension};
11pub use queries::{
12    check_policy_rate_limits, check_policy_rate_limits_for, get_all_rate_limits,
13    get_all_rate_limits_for, get_daily_usage, get_daily_usage_for, init_policy_rate_limits,
14    init_policy_rate_limits_for, record_policy_rate_limits, record_policy_rate_limits_for,
15    ActionUsage, DailyUsage,
16};
17pub use tracker::{
18    check_and_increment_rate_limit, check_and_increment_rate_limit_for, check_rate_limit,
19    check_rate_limit_for, increment_rate_limit, increment_rate_limit_for,
20};
21
22use super::DbPool;
23use crate::config::{IntervalsConfig, LimitsConfig};
24use crate::error::StorageError;
25
26use super::accounts::DEFAULT_ACCOUNT_ID;
27
28/// A rate limit entry tracking usage for a specific action type.
29#[derive(Debug, Clone, sqlx::FromRow, serde::Serialize)]
30pub struct RateLimit {
31    /// Action type: reply, tweet, thread, search, mention_check.
32    pub action_type: String,
33    /// Number of requests made in the current period.
34    pub request_count: i64,
35    /// ISO-8601 UTC timestamp when the current period started.
36    pub period_start: String,
37    /// Maximum requests allowed per period.
38    pub max_requests: i64,
39    /// Period length in seconds.
40    pub period_seconds: i64,
41}
42
43/// Initialize rate limit rows from configuration for a specific account.
44///
45/// Uses `INSERT OR IGNORE` so existing counters are preserved across restarts.
46/// Only inserts rows for action types that do not already exist.
47pub async fn init_rate_limits_for(
48    pool: &DbPool,
49    account_id: &str,
50    config: &LimitsConfig,
51    intervals: &IntervalsConfig,
52) -> Result<(), StorageError> {
53    // Suppress unused variable warning -- intervals is reserved for future per-interval limits
54    let _ = intervals;
55
56    let defaults: Vec<(&str, i64, i64)> = vec![
57        ("reply", i64::from(config.max_replies_per_day), 86400),
58        ("tweet", i64::from(config.max_tweets_per_day), 86400),
59        ("thread", i64::from(config.max_threads_per_week), 604800),
60        ("search", 300, 900),
61        ("mention_check", 180, 900),
62    ];
63
64    for (action_type, max_requests, period_seconds) in defaults {
65        sqlx::query(
66            "INSERT OR IGNORE INTO rate_limits \
67             (account_id, action_type, request_count, period_start, max_requests, period_seconds) \
68             VALUES (?, ?, 0, strftime('%Y-%m-%dT%H:%M:%SZ', 'now'), ?, ?)",
69        )
70        .bind(account_id)
71        .bind(action_type)
72        .bind(max_requests)
73        .bind(period_seconds)
74        .execute(pool)
75        .await
76        .map_err(|e| StorageError::Query { source: e })?;
77    }
78
79    Ok(())
80}
81
82/// Initialize rate limit rows from configuration.
83///
84/// Uses `INSERT OR IGNORE` so existing counters are preserved across restarts.
85/// Only inserts rows for action types that do not already exist.
86pub async fn init_rate_limits(
87    pool: &DbPool,
88    config: &LimitsConfig,
89    intervals: &IntervalsConfig,
90) -> Result<(), StorageError> {
91    init_rate_limits_for(pool, DEFAULT_ACCOUNT_ID, config, intervals).await
92}
93
94/// Initialize the MCP mutation rate limit row for a specific account.
95///
96/// Uses `INSERT OR IGNORE` so an existing counter is preserved across restarts.
97pub async fn init_mcp_rate_limit_for(
98    pool: &DbPool,
99    account_id: &str,
100    max_per_hour: u32,
101) -> Result<(), StorageError> {
102    sqlx::query(
103        "INSERT OR IGNORE INTO rate_limits \
104         (account_id, action_type, request_count, period_start, max_requests, period_seconds) \
105         VALUES (?, 'mcp_mutation', 0, strftime('%Y-%m-%dT%H:%M:%SZ', 'now'), ?, 3600)",
106    )
107    .bind(account_id)
108    .bind(i64::from(max_per_hour))
109    .execute(pool)
110    .await
111    .map_err(|e| StorageError::Query { source: e })?;
112
113    Ok(())
114}
115
116/// Initialize the MCP mutation rate limit row.
117///
118/// Uses `INSERT OR IGNORE` so an existing counter is preserved across restarts.
119pub async fn init_mcp_rate_limit(pool: &DbPool, max_per_hour: u32) -> Result<(), StorageError> {
120    init_mcp_rate_limit_for(pool, DEFAULT_ACCOUNT_ID, max_per_hour).await
121}
122
123#[cfg(test)]
124mod tests;
tuitbot_core/storage/rate_limits/mod.rs

tuitbot_core/storage/rate_limits/
mod.rs
