mcpr_core/protocol/schema_manager/

manager.rs

//! `SchemaManager` — top-level per-upstream view of an MCP server's schema.
//!
//! Callers feed schema-method responses in via [`SchemaManager::ingest`].
//! The manager handles pagination buffering, change detection (by content
//! hash), and version assignment, persisting new versions to a
//! [`SchemaStore`]. Query methods read back the latest merged payload
//! without re-hitting the store per item.

use std::sync::Arc;

use chrono::{DateTime, Utc};
use dashmap::DashMap;
use serde_json::Value;

use super::store::SchemaStore;
use super::version::{SchemaVersion, SchemaVersionId, hash_payload};
use crate::protocol::schema::{PageStatus, detect_page_status, merge_pages};

/// Per-method runtime state held in memory. Separate from the
/// `SchemaStore` because these fields serve the hot path — change
/// detection, pagination buffering, stale flag — and would be expensive
/// to read-through on every ingest.
#[derive(Default)]
struct MethodState {
    page_buffer: Vec<Value>,
    current_hash: Option<String>,
    next_version_number: u32,
    stale: bool,
    stale_since: Option<DateTime<Utc>>,
}

/// Top-level handle for one upstream MCP server's schema view.
///
/// Generic over the store backend; downstream typically uses
/// `SchemaManager<MemorySchemaStore>` for the OSS proxy and swaps in a
/// database-backed store for cloud deployments.
pub struct SchemaManager<S: SchemaStore> {
    upstream_id: String,
    store: S,
    state: Arc<DashMap<String, MethodState>>,
}

impl<S: SchemaStore> SchemaManager<S> {
    pub fn new(upstream_id: impl Into<String>, store: S) -> Self {
        Self {
            upstream_id: upstream_id.into(),
            store,
            state: Arc::new(DashMap::new()),
        }
    }

    pub fn upstream_id(&self) -> &str {
        &self.upstream_id
    }

    /// Seed the in-memory state for `method` from the store.
    ///
    /// Callers normally don't need to invoke this directly — `ingest`
    /// lazy-warms on the first call for a method. Exposed for explicit
    /// startup warm-up when desired.
    pub async fn warm(&self, method: &str) {
        let latest = self
            .store
            .latest_version_for_method(&self.upstream_id, method)
            .await;
        if let Some(latest) = latest {
            let mut entry = self.state.entry(method.to_string()).or_default();
            if entry.current_hash.is_none() {
                entry.current_hash = Some(latest.content_hash.clone());
                entry.next_version_number = latest.version + 1;
            }
        }
    }

    /// Feed a schema-method response through the manager.
    ///
    /// Returns `Some(version)` when a new `SchemaVersion` was created
    /// (pagination complete AND content differs from the current
    /// version). Returns `None` when:
    ///
    /// - The response is not a complete page (still buffering).
    /// - The content hash matches the current version.
    /// - The response has no `result` field.
    pub async fn ingest(
        &self,
        method: &str,
        request_body: &Value,
        response_body: &Value,
    ) -> Option<SchemaVersion> {
        let result = response_body.get("result")?;
        let status = detect_page_status(request_body, response_body);

        let merged = {
            let mut entry = self.state.entry(method.to_string()).or_default();
            entry.page_buffer.push(result.clone());
            match status {
                PageStatus::Complete | PageStatus::LastPage => {
                    let pages = std::mem::take(&mut entry.page_buffer);
                    merge_pages(method, &pages)
                        .unwrap_or_else(|| pages.into_iter().next().unwrap_or(Value::Null))
                }
                PageStatus::FirstPage | PageStatus::MiddlePage => return None,
            }
        };

        let hash = hash_payload(&merged);

        let needs_warm = self
            .state
            .get(method)
            .map(|e| e.current_hash.is_none() && e.next_version_number == 0)
            .unwrap_or(true);
        if needs_warm {
            self.warm(method).await;
        }

        let (same, version_number) = {
            let mut entry = self.state.entry(method.to_string()).or_default();
            if entry.current_hash.as_deref() == Some(hash.as_str()) {
                (true, 0)
            } else {
                let num = entry.next_version_number.max(1);
                entry.current_hash = Some(hash.clone());
                entry.next_version_number = num.saturating_add(1);
                entry.stale = false;
                entry.stale_since = None;
                (false, num)
            }
        };

        if same {
            return None;
        }

        let id = SchemaVersionId(hash.chars().take(16).collect());
        let version = SchemaVersion {
            id,
            upstream_id: self.upstream_id.clone(),
            method: method.to_string(),
            version: version_number,
            payload: Arc::new(merged),
            content_hash: hash,
            captured_at: Utc::now(),
        };
        Some(self.store.put_version(version).await)
    }

    /// Latest stored version for `method`, or `None` if nothing has
    /// been ingested yet.
    pub async fn latest(&self, method: &str) -> Option<SchemaVersion> {
        self.store
            .latest_version_for_method(&self.upstream_id, method)
            .await
    }

    pub async fn list_tools(&self) -> Vec<Value> {
        self.list_items("tools/list", "tools").await
    }

    pub async fn list_resources(&self) -> Vec<Value> {
        self.list_items("resources/list", "resources").await
    }

    pub async fn list_resource_templates(&self) -> Vec<Value> {
        self.list_items("resources/templates/list", "resourceTemplates")
            .await
    }

    pub async fn list_prompts(&self) -> Vec<Value> {
        self.list_items("prompts/list", "prompts").await
    }

    pub async fn get_tool(&self, name: &str) -> Option<Value> {
        self.find_item_by_field("tools/list", "tools", "name", name)
            .await
    }

    pub async fn get_resource(&self, uri: &str) -> Option<Value> {
        self.find_item_by_field("resources/list", "resources", "uri", uri)
            .await
    }

    pub async fn get_prompt(&self, name: &str) -> Option<Value> {
        self.find_item_by_field("prompts/list", "prompts", "name", name)
            .await
    }

    /// Mark the current version for `method` as stale. Idempotent.
    ///
    /// Sync on purpose — the stale flag is used by the hot request
    /// path (observing `notifications/tools/list_changed`) where a
    /// round-trip to async code would be overkill.
    pub fn mark_stale(&self, method: &str) {
        let mut entry = self.state.entry(method.to_string()).or_default();
        if !entry.stale {
            entry.stale = true;
            entry.stale_since = Some(Utc::now());
        }
    }

    pub fn is_stale(&self, method: &str) -> bool {
        self.state.get(method).map(|e| e.stale).unwrap_or(false)
    }

    pub fn stale_since(&self, method: &str) -> Option<DateTime<Utc>> {
        self.state.get(method).and_then(|e| e.stale_since)
    }

    // ── internals ──

    async fn list_items(&self, method: &str, array_key: &str) -> Vec<Value> {
        let Some(latest) = self.latest(method).await else {
            return Vec::new();
        };
        latest
            .payload
            .get(array_key)
            .and_then(|v| v.as_array())
            .cloned()
            .unwrap_or_default()
    }

    async fn find_item_by_field(
        &self,
        method: &str,
        array_key: &str,
        field: &str,
        needle: &str,
    ) -> Option<Value> {
        let latest = self.latest(method).await?;
        let arr = latest.payload.get(array_key).and_then(|v| v.as_array())?;
        arr.iter()
            .find(|item| item.get(field).and_then(|v| v.as_str()) == Some(needle))
            .cloned()
    }
}

#[cfg(test)]
#[allow(non_snake_case)]
mod tests {
    use super::*;
    use crate::protocol::schema_manager::store::MemorySchemaStore;
    use serde_json::json;

    fn manager() -> SchemaManager<MemorySchemaStore> {
        SchemaManager::new("proxy-1", MemorySchemaStore::new())
    }

    fn tools_list_req(cursor: Option<&str>) -> Value {
        match cursor {
            Some(c) => {
                json!({"jsonrpc": "2.0", "id": 1, "method": "tools/list", "params": {"cursor": c}})
            }
            None => json!({"jsonrpc": "2.0", "id": 1, "method": "tools/list"}),
        }
    }

    fn tools_list_resp(tools: Value, next_cursor: Option<&str>) -> Value {
        let mut result = json!({"tools": tools});
        if let Some(c) = next_cursor {
            result["nextCursor"] = json!(c);
        }
        json!({"jsonrpc": "2.0", "id": 1, "result": result})
    }

    #[tokio::test]
    async fn ingest__complete_page_creates_version_one() {
        let m = manager();
        let req = tools_list_req(None);
        let resp = tools_list_resp(json!([{"name": "search"}]), None);
        let v = m.ingest("tools/list", &req, &resp).await.unwrap();
        assert_eq!(v.version, 1);
        assert_eq!(v.method, "tools/list");
        assert_eq!(v.upstream_id, "proxy-1");
    }

    #[tokio::test]
    async fn ingest__first_page_buffers_returns_none() {
        let m = manager();
        let req = tools_list_req(None);
        let resp = tools_list_resp(json!([{"name": "a"}]), Some("cur1"));
        assert!(m.ingest("tools/list", &req, &resp).await.is_none());
    }

    #[tokio::test]
    async fn ingest__first_middle_last_chain_merges_once() {
        let m = manager();

        let r1 = tools_list_resp(json!([{"name": "a"}]), Some("c1"));
        assert!(
            m.ingest("tools/list", &tools_list_req(None), &r1)
                .await
                .is_none()
        );

        let r2 = tools_list_resp(json!([{"name": "b"}]), Some("c2"));
        assert!(
            m.ingest("tools/list", &tools_list_req(Some("c1")), &r2)
                .await
                .is_none()
        );

        let r3 = tools_list_resp(json!([{"name": "c"}]), None);
        let v = m
            .ingest("tools/list", &tools_list_req(Some("c2")), &r3)
            .await
            .unwrap();

        let names: Vec<&str> = v.payload["tools"]
            .as_array()
            .unwrap()
            .iter()
            .map(|t| t["name"].as_str().unwrap())
            .collect();
        assert_eq!(names, vec!["a", "b", "c"]);
        assert_eq!(v.version, 1);
    }

    #[tokio::test]
    async fn ingest__unchanged_payload_returns_none() {
        let m = manager();
        let req = tools_list_req(None);
        let resp = tools_list_resp(json!([{"name": "a"}]), None);
        m.ingest("tools/list", &req, &resp).await.unwrap();
        assert!(m.ingest("tools/list", &req, &resp).await.is_none());
    }

    #[tokio::test]
    async fn ingest__changed_payload_increments_version() {
        let m = manager();
        let req = tools_list_req(None);
        let r1 = tools_list_resp(json!([{"name": "a"}]), None);
        let v1 = m.ingest("tools/list", &req, &r1).await.unwrap();
        assert_eq!(v1.version, 1);

        let r2 = tools_list_resp(json!([{"name": "a"}, {"name": "b"}]), None);
        let v2 = m.ingest("tools/list", &req, &r2).await.unwrap();
        assert_eq!(v2.version, 2);
    }

    #[tokio::test]
    async fn ingest__clears_stale_on_new_version() {
        let m = manager();
        let req = tools_list_req(None);
        let r1 = tools_list_resp(json!([{"name": "a"}]), None);
        m.ingest("tools/list", &req, &r1).await.unwrap();

        m.mark_stale("tools/list");
        assert!(m.is_stale("tools/list"));

        let r2 = tools_list_resp(json!([{"name": "a"}, {"name": "b"}]), None);
        m.ingest("tools/list", &req, &r2).await.unwrap();
        assert!(!m.is_stale("tools/list"));
    }

    #[tokio::test]
    async fn ingest__no_result_returns_none() {
        let m = manager();
        let req = tools_list_req(None);
        let err_resp =
            json!({"jsonrpc": "2.0", "id": 1, "error": {"code": -32603, "message": "x"}});
        assert!(m.ingest("tools/list", &req, &err_resp).await.is_none());
    }

    #[tokio::test]
    async fn mark_stale__and_is_stale_idempotent() {
        let m = manager();
        assert!(!m.is_stale("tools/list"));
        m.mark_stale("tools/list");
        let first = m.stale_since("tools/list");
        m.mark_stale("tools/list");
        let second = m.stale_since("tools/list");
        assert!(m.is_stale("tools/list"));
        assert_eq!(first, second);
    }

    #[tokio::test]
    async fn list_tools__empty_when_no_version() {
        let m = manager();
        assert!(m.list_tools().await.is_empty());
    }

    #[tokio::test]
    async fn list_tools__returns_items_from_latest() {
        let m = manager();
        let req = tools_list_req(None);
        let resp = tools_list_resp(json!([{"name": "a"}, {"name": "b"}]), None);
        m.ingest("tools/list", &req, &resp).await.unwrap();

        let tools = m.list_tools().await;
        assert_eq!(tools.len(), 2);
        assert_eq!(tools[0]["name"], "a");
        assert_eq!(tools[1]["name"], "b");
    }

    #[tokio::test]
    async fn get_tool__by_name_hit_and_miss() {
        let m = manager();
        let req = tools_list_req(None);
        let resp = tools_list_resp(json!([{"name": "search", "description": "find"}]), None);
        m.ingest("tools/list", &req, &resp).await.unwrap();

        let hit = m.get_tool("search").await.unwrap();
        assert_eq!(hit["description"], "find");
        assert!(m.get_tool("missing").await.is_none());
    }

    #[tokio::test]
    async fn get_resource__by_uri() {
        let m = manager();
        let req = json!({"jsonrpc": "2.0", "id": 1, "method": "resources/list"});
        let resp = json!({
            "jsonrpc": "2.0", "id": 1,
            "result": {"resources": [{"uri": "file://a", "name": "A"}]}
        });
        m.ingest("resources/list", &req, &resp).await.unwrap();
        let r = m.get_resource("file://a").await.unwrap();
        assert_eq!(r["name"], "A");
    }

    #[tokio::test]
    async fn warm__seeds_counter_from_store() {
        let store = MemorySchemaStore::new();
        let pre = SchemaVersion {
            id: SchemaVersionId("abc".to_string()),
            upstream_id: "proxy-1".to_string(),
            method: "tools/list".to_string(),
            version: 5,
            payload: Arc::new(json!({"tools": [{"name": "x"}]})),
            content_hash: "prior-hash".to_string(),
            captured_at: Utc::now(),
        };
        store.put_version(pre).await;

        let m = SchemaManager::new("proxy-1", store);
        let req = tools_list_req(None);
        let resp = tools_list_resp(json!([{"name": "y"}]), None);
        let v = m.ingest("tools/list", &req, &resp).await.unwrap();
        assert_eq!(v.version, 6);
    }

    #[tokio::test]
    async fn latest__returns_current_version() {
        let m = manager();
        let req = tools_list_req(None);
        let resp = tools_list_resp(json!([{"name": "a"}]), None);
        m.ingest("tools/list", &req, &resp).await.unwrap();
        let latest = m.latest("tools/list").await.unwrap();
        assert_eq!(latest.version, 1);
    }

    #[tokio::test]
    async fn list_resource_templates__walks_template_key() {
        let m = manager();
        let req = json!({"jsonrpc": "2.0", "id": 1, "method": "resources/templates/list"});
        let resp = json!({
            "jsonrpc": "2.0", "id": 1,
            "result": {"resourceTemplates": [{"uriTemplate": "file://{id}", "name": "f"}]}
        });
        m.ingest("resources/templates/list", &req, &resp)
            .await
            .unwrap();
        let items = m.list_resource_templates().await;
        assert_eq!(items.len(), 1);
        assert_eq!(items[0]["name"], "f");
    }

    #[tokio::test]
    async fn upstream_id__accessor() {
        let m = manager();
        assert_eq!(m.upstream_id(), "proxy-1");
    }
}