tmcp 0.4.0

use std::collections::HashMap;

use serde::{Deserialize, Serialize};
use serde_json::Value;

/// Capabilities a client may support. Known capabilities are defined here,
/// in this schema, but this is not a closed set: any client can define its own,
/// additional capabilities.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ClientCapabilities {
    /// Experimental, non-standard capabilities that the client supports.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub experimental: Option<HashMap<String, Value>>,
    /// Present if the client supports listing roots.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub roots: Option<RootsCapability>,
    /// Present if the client supports sampling from an LLM.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub sampling: Option<SamplingCapability>,
    /// Present if the client supports elicitation from the server.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub elicitation: Option<ElicitationCapability>,
    /// Present if the client supports task-augmented requests.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tasks: Option<ClientTasksCapability>,
}

impl ClientCapabilities {
    /// Add an experimental capability
    pub fn with_experimental_capability(mut self, key: impl Into<String>, value: Value) -> Self {
        self.experimental
            .get_or_insert_with(HashMap::new)
            .insert(key.into(), value);
        self
    }

    /// Enable roots capability.
    ///
    /// list_changed indicates whether the client supports notifications for changes to the roots
    /// list.
    pub fn with_roots_capability(mut self, list_changed: bool) -> Self {
        self.roots = Some(RootsCapability {
            list_changed: Some(list_changed),
        });
        self
    }

    /// Enable sampling capability
    pub fn with_sampling(mut self) -> Self {
        self.sampling = Some(SamplingCapability::default());
        self
    }

    /// Enable elicitation capability.
    pub fn with_elicitation(mut self) -> Self {
        self.elicitation = Some(ElicitationCapability::default());
        self
    }

    /// Enable form-based elicitation capability.
    pub fn with_elicitation_form(mut self) -> Self {
        self.elicitation
            .get_or_insert_with(ElicitationCapability::default)
            .form = Some(HashMap::new());
        self
    }

    /// Enable URL-based elicitation capability.
    pub fn with_elicitation_url(mut self) -> Self {
        self.elicitation
            .get_or_insert_with(ElicitationCapability::default)
            .url = Some(HashMap::new());
        self
    }

    /// Enable task listing capability.
    pub fn with_tasks_list(mut self) -> Self {
        self.tasks
            .get_or_insert_with(ClientTasksCapability::default)
            .list = Some(HashMap::new());
        self
    }

    /// Enable task cancellation capability.
    pub fn with_tasks_cancel(mut self) -> Self {
        self.tasks
            .get_or_insert_with(ClientTasksCapability::default)
            .cancel = Some(HashMap::new());
        self
    }

    /// Enable task-augmented sampling/createMessage capability.
    pub fn with_task_sampling_create_message(mut self) -> Self {
        self.tasks
            .get_or_insert_with(ClientTasksCapability::default)
            .requests
            .get_or_insert_with(ClientTaskRequestsCapability::default)
            .sampling
            .get_or_insert_with(ClientTaskSamplingCapability::default)
            .create_message = Some(HashMap::new());
        self
    }

    /// Enable task-augmented elicitation/create capability.
    pub fn with_task_elicitation_create(mut self) -> Self {
        self.tasks
            .get_or_insert_with(ClientTasksCapability::default)
            .requests
            .get_or_insert_with(ClientTaskRequestsCapability::default)
            .elicitation
            .get_or_insert_with(ClientTaskElicitationCapability::default)
            .create = Some(HashMap::new());
        self
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RootsCapability {
    /// Whether the client supports notifications for changes to the roots list.
    #[serde(rename = "listChanged", skip_serializing_if = "Option::is_none")]
    pub list_changed: Option<bool>,
}

/// Capabilities that a server may support. Known capabilities are defined here,
/// in this schema, but this is not a closed set: any server can define its own,
/// additional capabilities.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ServerCapabilities {
    /// Experimental, non-standard capabilities that the server supports.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub experimental: Option<HashMap<String, Value>>,
    /// Present if the server supports sending log messages to the client.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub logging: Option<Value>,
    /// Present if the server supports argument autocompletion suggestions.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub completions: Option<Value>,
    /// Present if the server offers any prompt templates.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub prompts: Option<PromptsCapability>,
    /// Present if the server offers any resources to read.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub resources: Option<ResourcesCapability>,
    /// Present if the server offers any tools to call.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tools: Option<ToolsCapability>,
    /// Present if the server supports task-augmented requests.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tasks: Option<ServerTasksCapability>,
}

impl ServerCapabilities {
    /// Enable experimental capabilities
    pub fn with_experimental(mut self, experimental: HashMap<String, Value>) -> Self {
        self.experimental = Some(experimental);
        self
    }

    /// Enable logging capability
    pub fn with_logging(mut self) -> Self {
        self.logging = Some(Value::Object(serde_json::Map::new()));
        self
    }

    /// Enable completions capability
    pub fn with_completions(mut self) -> Self {
        self.completions = Some(Value::Object(serde_json::Map::new()));
        self
    }

    /// Enable prompts capability with optional list_changed support
    pub fn with_prompts(mut self, list_changed: Option<bool>) -> Self {
        self.prompts = Some(PromptsCapability { list_changed });
        self
    }

    /// Enable resources capability with optional subscribe and list_changed support
    pub fn with_resources(mut self, subscribe: Option<bool>, list_changed: Option<bool>) -> Self {
        self.resources = Some(ResourcesCapability {
            subscribe,
            list_changed,
        });
        self
    }

    /// Enable tools capability with optional list_changed support
    pub fn with_tools(mut self, list_changed: Option<bool>) -> Self {
        self.tools = Some(ToolsCapability { list_changed });
        self
    }

    /// Enable task listing capability.
    pub fn with_tasks_list(mut self) -> Self {
        self.tasks
            .get_or_insert_with(ServerTasksCapability::default)
            .list = Some(HashMap::new());
        self
    }

    /// Enable task cancellation capability.
    pub fn with_tasks_cancel(mut self) -> Self {
        self.tasks
            .get_or_insert_with(ServerTasksCapability::default)
            .cancel = Some(HashMap::new());
        self
    }

    /// Enable task-augmented tools/call capability.
    pub fn with_task_tools_call(mut self) -> Self {
        self.tasks
            .get_or_insert_with(ServerTasksCapability::default)
            .requests
            .get_or_insert_with(ServerTaskRequestsCapability::default)
            .tools
            .get_or_insert_with(ServerTaskToolsCapability::default)
            .call = Some(HashMap::new());
        self
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct SamplingCapability {
    /// Whether the client supports context inclusion via includeContext parameter.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub context: Option<HashMap<String, Value>>,
    /// Whether the client supports tool use via tools and toolChoice parameters.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tools: Option<HashMap<String, Value>>,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ElicitationCapability {
    /// Support for form-based elicitation.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub form: Option<HashMap<String, Value>>,
    /// Support for URL-based elicitation.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub url: Option<HashMap<String, Value>>,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ClientTasksCapability {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub list: Option<HashMap<String, Value>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cancel: Option<HashMap<String, Value>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub requests: Option<ClientTaskRequestsCapability>,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ClientTaskRequestsCapability {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub sampling: Option<ClientTaskSamplingCapability>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub elicitation: Option<ClientTaskElicitationCapability>,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ClientTaskSamplingCapability {
    #[serde(rename = "createMessage", skip_serializing_if = "Option::is_none")]
    pub create_message: Option<HashMap<String, Value>>,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ClientTaskElicitationCapability {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub create: Option<HashMap<String, Value>>,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ServerTasksCapability {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub list: Option<HashMap<String, Value>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cancel: Option<HashMap<String, Value>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub requests: Option<ServerTaskRequestsCapability>,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ServerTaskRequestsCapability {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tools: Option<ServerTaskToolsCapability>,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ServerTaskToolsCapability {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub call: Option<HashMap<String, Value>>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PromptsCapability {
    /// Whether this server supports notifications for changes to the prompt
    /// list.
    #[serde(rename = "listChanged", skip_serializing_if = "Option::is_none")]
    pub list_changed: Option<bool>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResourcesCapability {
    /// Whether this server supports subscribing to resource updates.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub subscribe: Option<bool>,
    /// Whether this server supports notifications for changes to the resource
    /// list.
    #[serde(rename = "listChanged", skip_serializing_if = "Option::is_none")]
    pub list_changed: Option<bool>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolsCapability {
    /// Whether this server supports notifications for changes to the tool list.
    #[serde(rename = "listChanged", skip_serializing_if = "Option::is_none")]
    pub list_changed: Option<bool>,
}