heptagent-memory-tool-backend 0.1.0

// SPDX-License-Identifier: MIT
// Copyright (c) 2026 heptagent core team

//! [`QuestRateLimiter`] — D-C-03 Worker active `memory_20250818` tool-call rate limit.
//!
//! 20 calls/quest hard cap (T-14-00-03 DoS mitigation).
//! `AtomicU32` provides lock-free try_consume + reset_for_new_quest per NP42.

use std::sync::atomic::{AtomicU32, Ordering};

use crate::error::MemoryError;

/// Rate limiter for Worker `memory_20250818` tool-calls per quest.
///
/// **D-C-03 locked decision**: 20 calls/quest hard cap.
/// After 20 calls, `try_consume` returns `Err(MemoryError::RateLimit)`.
/// Caller (Coord) calls `reset_for_new_quest` at quest boundary.
///
/// **NP42** — `QuestRateLimiter` is NOT a trait (single impl, no object-safety needed).
pub struct QuestRateLimiter {
    calls_this_quest: AtomicU32,
    max_calls: u32,
}

impl QuestRateLimiter {
    /// Create a new rate limiter with `max_calls` cap.
    ///
    /// Standard construction: `QuestRateLimiter::new(20)` per D-C-03.
    pub fn new(max_calls: u32) -> Self {
        Self {
            calls_this_quest: AtomicU32::new(0),
            max_calls,
        }
    }

    /// Attempt to consume one call slot.
    ///
    /// Returns `Ok(())` if under cap, `Err(MemoryError::RateLimit)` on 21st+ call.
    ///
    /// Uses `Ordering::Relaxed` — correctness is approximate (per-quest budget,
    /// not a strict security boundary). Production K-08 may upgrade to `SeqCst`
    /// if multi-threaded dispatcher dispatch is introduced.
    pub fn try_consume(&self) -> Result<(), MemoryError> {
        let prev = self.calls_this_quest.fetch_add(1, Ordering::Relaxed);
        if prev >= self.max_calls {
            Err(MemoryError::RateLimit {
                limit: self.max_calls,
            })
        } else {
            Ok(())
        }
    }

    /// Reset call counter to 0 at quest boundary.
    pub fn reset_for_new_quest(&self) {
        self.calls_this_quest.store(0, Ordering::Relaxed);
    }

    /// Current call count (for observability).
    pub fn current_calls(&self) -> u32 {
        self.calls_this_quest.load(Ordering::Relaxed)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn rate_limiter_allows_20_calls() {
        let rl = QuestRateLimiter::new(20);
        for _ in 0..20 {
            assert!(rl.try_consume().is_ok());
        }
    }

    #[test]
    fn rate_limiter_rejects_21st_call() {
        let rl = QuestRateLimiter::new(20);
        for _ in 0..20 {
            let _ = rl.try_consume();
        }
        assert!(matches!(
            rl.try_consume(),
            Err(MemoryError::RateLimit { limit: 20 })
        ));
    }

    #[test]
    fn rate_limiter_reset_allows_again() {
        let rl = QuestRateLimiter::new(20);
        for _ in 0..20 {
            let _ = rl.try_consume();
        }
        rl.reset_for_new_quest();
        assert!(rl.try_consume().is_ok());
    }
}