~ knotless Search/
↑↓ move↵ open⇧↵ all results

johnandersen777.bsky.social/asyncai

0
0
main
164 lines · 5.7 KB Raw History
use rmcp::handler::server::router::tool::ToolRouter;
use rmcp::handler::server::wrapper::Parameters;
use std::sync::Arc;

use rmcp::model::{
    CallToolResult, ListResourceTemplatesResult, PaginatedRequestParams, ReadResourceRequestParams,
    ReadResourceResponse, ServerConfig, SubscriptionFilter,
};
use rmcp::service::{RequestContext, SubscriptionContext};
use rmcp::transport::stdio;
use rmcp::{tool, tool_handler, tool_router, ServerHandler, ServiceExt};
use rmcp::{ErrorData, RoleServer};

use super::service::KcpServiceHandle;
use super::watch::{self, WatchHub};
use super::{spec, tools};

#[derive(Clone)]
pub struct ReadwriteServer {
    tool_router: ToolRouter<Self>,
    service: KcpServiceHandle,
    watches: Arc<WatchHub>,
}

impl ReadwriteServer {
    pub fn new(service: KcpServiceHandle) -> Self {
        Self {
            tool_router: Self::tool_router(),
            service,
            watches: Arc::default(),
        }
    }
}

#[tool_router]
impl ReadwriteServer {
    #[tool(description = spec::LIST_KCPS, output_schema = spec::kcp_report(), annotations(read_only_hint = true))]
    pub async fn list_kcps(
        &self,
        Parameters(request): Parameters<tools::ListKcpsRequest>,
    ) -> CallToolResult {
        tools::list_kcps(&self.service, request).await
    }

    #[tool(description = spec::LIST_WORKSPACES, output_schema = spec::workspaces(), annotations(read_only_hint = true))]
    pub async fn list_workspaces(
        &self,
        Parameters(request): Parameters<tools::ListWorkspacesRequest>,
    ) -> CallToolResult {
        tools::list_workspaces(&self.service, request).await
    }

    #[tool(description = spec::GET_API_RESOURCES, output_schema = spec::api_resources(), annotations(read_only_hint = true))]
    pub async fn get_api_resources(
        &self,
        Parameters(request): Parameters<tools::KcpWorkspaceRequest>,
    ) -> CallToolResult {
        tools::get_api_resources(&self.service, request).await
    }

    #[tool(description = spec::LIST_OBJECTS, output_schema = spec::kube_list(), annotations(read_only_hint = true))]
    pub async fn list_objects(
        &self,
        Parameters(request): Parameters<tools::ListObjectsRequest>,
    ) -> CallToolResult {
        tools::list_objects(&self.service, request).await
    }

    #[tool(description = spec::GET_OBJECT, output_schema = spec::kube_object(), annotations(read_only_hint = true))]
    pub async fn get_object(
        &self,
        Parameters(request): Parameters<tools::GetObjectRequest>,
    ) -> CallToolResult {
        tools::get_object(&self.service, request).await
    }

    #[tool(description = spec::GET_LOGS, annotations(read_only_hint = true))]
    pub async fn get_logs(
        &self,
        Parameters(request): Parameters<tools::GetLogsRequest>,
    ) -> CallToolResult {
        tools::get_logs(&self.service, request).await
    }

    #[tool(description = spec::GET_EVENTS, output_schema = spec::event_list(), annotations(read_only_hint = true))]
    pub async fn get_events(
        &self,
        Parameters(request): Parameters<tools::GetEventsRequest>,
    ) -> CallToolResult {
        tools::get_events(&self.service, request).await
    }

    #[tool(description = spec::APPLY_OBJECT, output_schema = spec::kube_object(), annotations(read_only_hint = false, destructive_hint = false))]
    pub async fn apply_object(
        &self,
        Parameters(request): Parameters<tools::ApplyObjectRequest>,
    ) -> CallToolResult {
        tools::apply_object(&self.service, request).await
    }

    #[tool(description = spec::DELETE_OBJECT, annotations(read_only_hint = false, destructive_hint = true))]
    pub async fn delete_object(
        &self,
        Parameters(request): Parameters<tools::DeleteObjectRequest>,
    ) -> CallToolResult {
        tools::delete_object(&self.service, request).await
    }

    #[tool(description = spec::PATCH_OBJECT, output_schema = spec::kube_object(), annotations(read_only_hint = false, destructive_hint = false))]
    pub async fn patch_object(
        &self,
        Parameters(request): Parameters<tools::PatchObjectRequest>,
    ) -> CallToolResult {
        tools::patch_object(&self.service, request).await
    }

    #[tool(description = spec::SCALE_OBJECT, output_schema = spec::kube_object(), annotations(read_only_hint = false, destructive_hint = false))]
    pub async fn scale_object(
        &self,
        Parameters(request): Parameters<tools::ScaleObjectRequest>,
    ) -> CallToolResult {
        tools::scale_object(&self.service, request).await
    }
}

#[tool_handler(router = self.tool_router)]
impl ServerHandler for ReadwriteServer {
    fn get_info(&self) -> ServerConfig {
        watch::server_info()
    }

    async fn list_resource_templates(
        &self,
        _request: Option<PaginatedRequestParams>,
        _context: RequestContext<RoleServer>,
    ) -> Result<ListResourceTemplatesResult, ErrorData> {
        Ok(watch::resource_templates())
    }

    async fn read_resource(
        &self,
        request: ReadResourceRequestParams,
        _context: RequestContext<RoleServer>,
    ) -> Result<ReadResourceResponse, ErrorData> {
        watch::read_resource(&self.watches, &self.service, &request.uri).await
    }

    fn accepted_subscription_filter(
        &self,
        requested: &SubscriptionFilter,
    ) -> Option<SubscriptionFilter> {
        watch::accepted_subscriptions(requested)
    }

    async fn listen(&self, context: SubscriptionContext) -> Result<(), ErrorData> {
        watch::listen(&self.watches, &self.service, context).await
    }
}

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