docs_mcp/tools/
mod.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use governor::{DefaultDirectRateLimiter, Quota, RateLimiter};
5use http::Extensions;
6use nonzero_ext::nonzero;
7use reqwest::Request;
8use reqwest_middleware::{Middleware, Next};
9
10use crate::cache::DiskCache;
11use crate::error::Result;
12use crate::sparse_index::{self, IndexLine};
13
14pub mod crate_list;
15pub mod crate_get;
16pub mod crate_readme_get;
17pub mod crate_docs_get;
18pub mod crate_item_list;
19pub mod crate_item_get;
20pub mod crate_impls_list;
21pub mod crate_versions_list;
22pub mod crate_version_get;
23pub mod crate_dependencies_list;
24pub mod crate_dependents_list;
25pub mod crate_downloads_get;
26
27/// Shared application state, held behind an Arc in the server.
28pub struct AppState {
29    pub client: reqwest_middleware::ClientWithMiddleware,
30    pub cache: DiskCache,
31}
32
33impl AppState {
34    pub async fn new() -> Result<Self> {
35        let mut headers = reqwest::header::HeaderMap::new();
36        headers.insert(
37            reqwest::header::USER_AGENT,
38            reqwest::header::HeaderValue::from_static(
39                "docs-mcp/0.1 (https://github.com/user/docs-mcp)",
40            ),
41        );
42
43        let http = reqwest::Client::builder()
44            .default_headers(headers)
45            .build()
46            .map_err(crate::error::DocsError::Http)?;
47
48        let rate_mw = RateLimitMiddleware::new();
49        let cache = DiskCache::new()?;
50
51        let client = reqwest_middleware::ClientBuilder::new(http)
52            .with(rate_mw)
53            .build();
54
55        Ok(Self { client, cache })
56    }
57
58    /// Resolve a version string: if None or "latest", look up the latest stable version.
59    pub async fn resolve_version(&self, name: &str, version: Option<&str>) -> Result<String> {
60        match version {
61            Some(v) if !v.is_empty() && v != "latest" => Ok(v.to_string()),
62            _ => {
63                let lines = sparse_index::fetch_index(name, &self.client, &self.cache).await?;
64                let latest = sparse_index::find_latest_stable(&lines)
65                    .ok_or_else(|| crate::error::DocsError::NoStableVersion(name.to_string()))?;
66                Ok(latest.vers.clone())
67            }
68        }
69    }
70
71    /// Fetch all index lines for a crate.
72    pub async fn fetch_index(&self, name: &str) -> Result<Vec<IndexLine>> {
73        sparse_index::fetch_index(name, &self.client, &self.cache).await
74    }
75}
76
77// ─── Rate limit middleware ─────────────────────────────────────────────────────
78
79pub struct RateLimitMiddleware {
80    limiter: Arc<DefaultDirectRateLimiter>,
81}
82
83impl RateLimitMiddleware {
84    pub fn new() -> Self {
85        let quota = Quota::per_second(nonzero!(1u32));
86        let limiter = Arc::new(RateLimiter::direct(quota));
87        Self { limiter }
88    }
89}
90
91#[async_trait]
92impl Middleware for RateLimitMiddleware {
93    async fn handle(
94        &self,
95        req: Request,
96        extensions: &mut Extensions,
97        next: Next<'_>,
98    ) -> reqwest_middleware::Result<reqwest::Response> {
99        // Only rate limit crates.io API calls (not sparse index or docs.rs)
100        if req.url().host_str() == Some("crates.io") {
101            self.limiter.until_ready().await;
102        }
103        next.run(req, extensions).await
104    }
105}
docs_mcp/tools/mod.rs

docs_mcp/tools/
mod.rs
