johnandersen777.bsky.social/asyncai

Switch repository

Loading…

use std::collections::BTreeMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;

use rmcp::handler::server::{router::tool::ToolRouter, wrapper::Parameters};
use rmcp::model::{
    CallToolResult, ContentBlock, JsonObject, ListResourcesResult, PaginatedRequestParams,
    ReadResourceRequestParams, ReadResourceResponse, ReadResourceResult, Resource,
    ResourceContents, ServerCapabilities, ServerConfig, SubscriptionFilter,
};
use rmcp::schemars::JsonSchema;
use rmcp::service::{RequestContext, SubscriptionContext};
use rmcp::transport::stdio;
use rmcp::{tool, tool_handler, tool_router, ErrorData, RoleServer, ServerHandler, ServiceExt};
use serde::{Deserialize, Serialize};
use serde_json::json;
use tokio::sync::broadcast;

/// State resource demonstrating the ADR 0005 watch contract.
pub const MEMO_URI: &str = "memo://notes";

#[derive(Clone)]
pub struct SimpleServer {
    tool_router: ToolRouter<Self>,
    memo: Arc<Mutex<BTreeMap<String, String>>>,
    changed: broadcast::Sender<()>,
}

impl SimpleServer {
    pub fn new() -> Self {
        Self {
            tool_router: Self::tool_router(),
            memo: Arc::default(),
            changed: broadcast::channel(64).0,
        }
    }

    fn memo_json(&self) -> String {
        json!(*self.memo.lock().expect("memo lock")).to_string()
    }
}

impl Default for SimpleServer {
    fn default() -> Self {
        Self::new()
    }
}

#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct HttpGetRequest {
    /// Absolute http(s) URL to fetch.
    pub url: String,
}

#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct MemoSetRequest {
    pub key: String,
    /// New value; omit to delete the key.
    pub value: Option<String>,
}

#[tool_router]
impl SimpleServer {
    #[tool(
        description = "Perform an HTTP GET request and return the full body. Non-2xx statuses are results, not errors; connection failures raise.\nExample: r = await http_get(url=\"https://example.com/data.json\"); json.loads(r[\"body\"]) if r[\"status\"] == 200 else r[\"status\"]",
        output_schema = http_get_schema(),
        annotations(read_only_hint = true)
    )]
    pub async fn http_get(
        &self,
        Parameters(request): Parameters<HttpGetRequest>,
    ) -> CallToolResult {
        let client = reqwest::Client::new();
        let response = match client
            .get(&request.url)
            .timeout(Duration::from_secs(15))
            .send()
            .await
        {
            Ok(response) => response,
            Err(error) => return failure(format!("http_get error: {error}")),
        };
        let status = response.status().as_u16();
        let content_type = response
            .headers()
            .get(reqwest::header::CONTENT_TYPE)
            .and_then(|value| value.to_str().ok())
            .unwrap_or("")
            .to_string();
        match response.text().await {
            Ok(body) => CallToolResult::structured(
                json!({"status": status, "content_type": content_type, "body": body}),
            ),
            Err(error) => failure(format!("http_get body read error: {error}")),
        }
    }

    #[tool(
        description = "Set (or delete, with no value) one note in the memo://notes resource. Watchers of memo://notes are notified.\nExample: await memo_set(key=\"todo\", value=\"ship 0005\")",
        annotations(read_only_hint = false, destructive_hint = false)
    )]
    pub async fn memo_set(
        &self,
        Parameters(request): Parameters<MemoSetRequest>,
    ) -> CallToolResult {
        {
            let mut memo = self.memo.lock().expect("memo lock");
            match &request.value {
                Some(value) => memo.insert(request.key.clone(), value.clone()),
                None => memo.remove(&request.key),
            };
        }
        let _ = self.changed.send(());
        CallToolResult::structured(
            json!({"notes": serde_json::from_str::<serde_json::Value>(&self.memo_json()).unwrap_or_default()}),
        )
    }
}

#[tool_handler(router = self.tool_router)]
impl ServerHandler for SimpleServer {
    fn get_info(&self) -> ServerConfig {
        ServerConfig::new(
            ServerCapabilities::builder()
                .enable_tools()
                .enable_resources()
                .enable_resources_subscribe()
                .build(),
        )
    }

    async fn list_resources(
        &self,
        _request: Option<PaginatedRequestParams>,
        _context: RequestContext<RoleServer>,
    ) -> Result<ListResourcesResult, ErrorData> {
        Ok(ListResourcesResult::with_all_items(vec![
            Resource::new(MEMO_URI, "notes")
                .with_description("Shared notes (JSON object). Watch it to hear about changes.\nExample: await watch_start(name=\"notes\", uri=\"memo://notes\")")
                .with_mime_type("application/json"),
        ]))
    }

    async fn read_resource(
        &self,
        request: ReadResourceRequestParams,
        _context: RequestContext<RoleServer>,
    ) -> Result<ReadResourceResponse, ErrorData> {
        if request.uri != MEMO_URI {
            return Err(ErrorData::resource_not_found(
                format!("unknown resource {}", request.uri),
                None,
            ));
        }
        Ok(ReadResourceResponse::Complete(ReadResourceResult::new(
            vec![ResourceContents::TextResourceContents {
                uri: MEMO_URI.to_string(),
                mime_type: Some("application/json".to_string()),
                text: self.memo_json(),
                meta: None,
            }],
        )))
    }

    fn accepted_subscription_filter(
        &self,
        requested: &SubscriptionFilter,
    ) -> Option<SubscriptionFilter> {
        let wants_memo = requested
            .resource_subscriptions
            .as_ref()
            .is_some_and(|uris| uris.iter().any(|uri| uri == MEMO_URI));
        Some(if wants_memo {
            SubscriptionFilter::builder()
                .resource_subscription(MEMO_URI)
                .build()
        } else {
            SubscriptionFilter::new()
        })
    }

    async fn listen(&self, context: SubscriptionContext) -> Result<(), ErrorData> {
        let mut changed = self.changed.subscribe();
        loop {
            tokio::select! {
                _ = context.cancelled() => return Ok(()),
                result = changed.recv() => match result {
                    Ok(()) | Err(broadcast::error::RecvError::Lagged(_)) => {
                        if context.sink().notify_resource_updated(MEMO_URI).await.is_err() {
                            return Ok(());
                        }
                    }
                    Err(broadcast::error::RecvError::Closed) => return Ok(()),
                },
            }
        }
    }
}

fn failure(message: String) -> CallToolResult {
    CallToolResult::error(vec![ContentBlock::text(message)])
}

fn http_get_schema() -> Arc<JsonObject> {
    match json!({
        "type": "object",
        "properties": {
            "status": {"type": "integer"},
            "content_type": {"type": "string"},
            "body": {"type": "string"},
        },
        "additionalProperties": false,
    }) {
        serde_json::Value::Object(object) => Arc::new(object),
        _ => unreachable!(),
    }
}

pub async fn serve_stdio() -> Result<(), Box<dyn std::error::Error>> {
    let service = SimpleServer::new().serve(stdio()).await?;
    service.waiting().await?;
    Ok(())
}