wowlab_sentinel/cron/

docs.rs

use std::time::Duration;

use async_trait::async_trait;
use futures::StreamExt;
use reqwest::Client;
use wowlab_types::constants::{
    BYTES_PER_MB_USIZE, HTTP_CONNECT_TIMEOUT_SECS, HTTP_REQUEST_TIMEOUT_SECS,
};

use super::CronJob;
use crate::state::{McpStateHandle, ServerState};

pub(crate) const CACHE_DOCS_INDEX: &str = "docs_index";
pub(crate) const CACHE_DOCS_FULL: &str = "docs_full";
const MAX_DOC_BYTES: usize = 2 * BYTES_PER_MB_USIZE;

#[derive(Debug)]
pub(crate) struct DocsCacheJob {
    schedule: String,
}

impl DocsCacheJob {
    pub(crate) fn new(schedule: &str) -> Self {
        Self {
            schedule: schedule.to_owned(),
        }
    }
}

#[async_trait]
impl CronJob for DocsCacheJob {
    fn name(&self) -> &'static str {
        "docs_cache"
    }

    fn schedule(&self) -> &str {
        &self.schedule
    }

    async fn run(&self, state: &ServerState) {
        refresh(&state.config.docs_base_url, &state.mcp).await;
    }
}

pub(crate) async fn refresh(base_url: &str, mcp: &McpStateHandle) {
    let client = match Client::builder()
        .connect_timeout(Duration::from_secs(HTTP_CONNECT_TIMEOUT_SECS))
        .timeout(Duration::from_secs(HTTP_REQUEST_TIMEOUT_SECS))
        .build()
    {
        Ok(client) => client,
        Err(error) => {
            tracing::warn!(%error, "Failed to configure docs HTTP client");

            return;
        }
    };

    refresh_with_client(base_url, mcp, &client, MAX_DOC_BYTES).await;
}

async fn refresh_with_client(
    base_url: &str,
    mcp: &McpStateHandle,
    client: &Client,
    max_bytes: usize,
) {
    let index_url = format!("{base_url}/llms.txt");
    let full_url = format!("{base_url}/llms-full.txt");
    let (index, full) = tokio::join!(
        fetch_text(client, &index_url, max_bytes),
        fetch_text(client, &full_url, max_bytes)
    );

    let mut c = mcp.docs_mut().await;

    if let Some(t) = index {
        c.set(CACHE_DOCS_INDEX, t);
    }

    if let Some(t) = full {
        c.set(CACHE_DOCS_FULL, t);
    }

    tracing::info!(
        index = c.get(CACHE_DOCS_INDEX).map_or(0, str::len),
        full = c.get(CACHE_DOCS_FULL).map_or(0, str::len),
        "Docs cache refreshed"
    );
}

async fn fetch_text(client: &Client, url: &str, max_bytes: usize) -> Option<String> {
    let response = match client.get(url).send().await {
        Ok(response) => response,
        Err(error) => {
            tracing::warn!(url, %error, "Failed to fetch docs");

            return None;
        }
    };

    if !response.status().is_success() {
        tracing::warn!(url, status = %response.status(), "Failed to fetch docs");

        return None;
    }

    if response
        .content_length()
        .is_some_and(|length| usize::try_from(length).map_or(true, |length| length > max_bytes))
    {
        tracing::warn!(url, max_bytes, "Docs response exceeds size limit");

        return None;
    }

    let mut body = Vec::with_capacity(
        response
            .content_length()
            .and_then(|length| usize::try_from(length).ok())
            .unwrap_or_default()
            .min(max_bytes),
    );
    let mut stream = response.bytes_stream();
    let mut read_error = None;
    let mut oversized = false;

    loop {
        let Some(chunk) = stream.next().await else {
            break;
        };
        let chunk = match chunk {
            Ok(chunk) => chunk,
            Err(error) => {
                read_error = Some(error);
                break;
            }
        };

        if chunk.len() > max_bytes.saturating_sub(body.len()) {
            oversized = true;
            break;
        }

        body.extend_from_slice(&chunk);
    }

    if let Some(error) = read_error {
        tracing::warn!(url, %error, "Failed to read docs response");

        return None;
    }

    if oversized {
        tracing::warn!(url, max_bytes, "Docs response exceeds size limit");

        return None;
    }

    match String::from_utf8(body) {
        Ok(body) => Some(body),
        Err(error) => {
            tracing::warn!(url, %error, "Docs response is not valid UTF-8");

            None
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use googletest::prelude::*;
    use wiremock::{Mock, MockServer, ResponseTemplate, matchers::path};

    use super::*;
    use crate::state::McpState;

    #[gtest]
    #[tokio::test]
    async fn oversized_response_does_not_replace_cached_docs() -> Result<()> {
        let server = MockServer::start().await;

        Mock::given(path("/llms.txt"))
            .respond_with(ResponseTemplate::new(200).set_body_string("too large"))
            .mount(&server)
            .await;

        Mock::given(path("/llms-full.txt"))
            .respond_with(ResponseTemplate::new(200).set_body_string("new full"))
            .mount(&server)
            .await;

        let mcp = Arc::new(McpState::new()?);

        mcp.docs_mut()
            .await
            .set(CACHE_DOCS_INDEX, "old index".to_owned());

        refresh_with_client(&server.uri(), &mcp, &Client::new(), 8).await;

        let docs = mcp.docs().await;

        verify_that!(docs.get(CACHE_DOCS_INDEX), eq(Some("old index")))?;
        verify_that!(docs.get(CACHE_DOCS_FULL), eq(Some("new full")))?;

        Ok(())
    }

    #[gtest]
    #[tokio::test]
    async fn request_timeout_is_nonfatal() -> Result<()> {
        let server = MockServer::start().await;

        Mock::given(path("/llms.txt"))
            .respond_with(
                ResponseTemplate::new(200)
                    .set_delay(Duration::from_millis(200))
                    .set_body_string("late"),
            )
            .mount(&server)
            .await;
        let client = Client::builder()
            .connect_timeout(Duration::from_millis(25))
            .timeout(Duration::from_millis(25))
            .build()?;

        let result = fetch_text(&client, &format!("{}/llms.txt", server.uri()), 1024).await;

        verify_true!(result.is_none())?;

        Ok(())
    }
}