magi-code 0.104.0

//! Reader-owned notification flags and bounded, capability-free server replies.
use super::jsonrpc::{JsonRpcNotification, JsonRpcRequest, METHOD_NOT_FOUND};
use serde_json::{Value, json};
use std::{
    sync::{
        Mutex,
        atomic::{AtomicBool, Ordering},
    },
    time::{Duration, Instant},
};

#[derive(Debug)]
pub(crate) struct ServerMessages {
    list_changed_enabled: AtomicBool,
    tools_stale: AtomicBool,
    reply_window: Mutex<(Instant, usize)>,
}

impl Default for ServerMessages {
    fn default() -> Self {
        Self {
            list_changed_enabled: AtomicBool::new(false),
            tools_stale: AtomicBool::new(false),
            reply_window: Mutex::new((Instant::now(), 0)),
        }
    }
}

impl ServerMessages {
    pub(crate) fn enable_list_changed(&self, enabled: bool) {
        self.list_changed_enabled.store(enabled, Ordering::Release);
    }

    pub(crate) fn notification(&self, notification: &JsonRpcNotification) {
        if notification.method == "notifications/tools/list_changed"
            && self.list_changed_enabled.load(Ordering::Acquire)
        {
            self.tools_stale.store(true, Ordering::Release);
        }
    }

    pub(crate) fn take_tools_stale(&self) -> bool {
        self.tools_stale.swap(false, Ordering::AcqRel)
    }

    pub(crate) fn reply(&self, request: JsonRpcRequest) -> Option<Value> {
        // One reader writes synchronously on stdio; HTTP additionally caps its queue at 32.
        // Rate limiting also bounds a server that sends requests faster than writes finish.
        let mut window = self.reply_window.lock().ok()?;
        if window.0.elapsed() >= Duration::from_secs(1) {
            *window = (Instant::now(), 0);
        }
        if window.1 >= 32 {
            return None;
        }
        window.1 += 1;
        Some(if request.method == "ping" {
            json!({"jsonrpc":"2.0", "id":request.id, "result":{}})
        } else {
            json!({"jsonrpc":"2.0", "id":request.id,
                "error":{"code":METHOD_NOT_FOUND, "message":"Method not found"}})
        })
    }
}