rmcp 2.1.0

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

use super::Meta;

/// Metadata for augmenting a request with task execution (spec `TaskMetadata`).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[non_exhaustive]
pub struct TaskMetadata {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub ttl: Option<u64>,
}

impl TaskMetadata {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_ttl(mut self, ttl: u64) -> Self {
        self.ttl = Some(ttl);
        self
    }
}

/// Metadata for associating messages with a task (spec `RelatedTaskMetadata`).
///
/// Carried in `_meta` under the key `"io.modelcontextprotocol/related-task"`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[non_exhaustive]
pub struct RelatedTaskMetadata {
    pub task_id: String,
}

impl RelatedTaskMetadata {
    pub fn new(task_id: impl Into<String>) -> Self {
        Self {
            task_id: task_id.into(),
        }
    }

    /// The well-known `_meta` key for related-task metadata.
    pub const META_KEY: &str = "io.modelcontextprotocol/related-task";
}

/// Canonical task lifecycle status as defined by SEP-1686.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[non_exhaustive]
pub enum TaskStatus {
    /// The receiver accepted the request and is currently working on it.
    #[default]
    Working,
    /// The receiver requires additional input before work can continue.
    InputRequired,
    /// The underlying operation completed successfully and the result is ready.
    Completed,
    /// The underlying operation failed and will not continue.
    Failed,
    /// The task was cancelled and will not continue processing.
    Cancelled,
}

/// Primary Task object that surfaces metadata during the task lifecycle.
///
/// Per spec, `lastUpdatedAt` and `ttl` are required fields.
/// `ttl` is nullable (`null` means unlimited retention).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[non_exhaustive]
pub struct Task {
    /// Unique task identifier generated by the receiver.
    pub task_id: String,
    /// Current lifecycle status (see [`TaskStatus`]).
    pub status: TaskStatus,
    /// Optional human-readable status message for UI surfaces.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub status_message: Option<String>,
    /// ISO-8601 creation timestamp.
    pub created_at: String,
    /// ISO-8601 timestamp for the most recent status change.
    pub last_updated_at: String,
    /// Retention window in milliseconds that the receiver agreed to honor.
    /// `None` (serialized as `null`) means unlimited retention.
    pub ttl: Option<u64>,
    /// Suggested polling interval (milliseconds).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub poll_interval: Option<u64>,
}

impl Task {
    /// Create a new Task with required fields.
    pub fn new(
        task_id: String,
        status: TaskStatus,
        created_at: String,
        last_updated_at: String,
    ) -> Self {
        Self {
            task_id,
            status,
            status_message: None,
            created_at,
            last_updated_at,
            ttl: None,
            poll_interval: None,
        }
    }

    /// Set the status message.
    pub fn with_status_message(mut self, status_message: impl Into<String>) -> Self {
        self.status_message = Some(status_message.into());
        self
    }

    /// Set the TTL in milliseconds. `None` means unlimited retention.
    pub fn with_ttl(mut self, ttl: u64) -> Self {
        self.ttl = Some(ttl);
        self
    }

    /// Set the poll interval in milliseconds.
    pub fn with_poll_interval(mut self, poll_interval: u64) -> Self {
        self.poll_interval = Some(poll_interval);
        self
    }
}

/// Wrapper returned by task-augmented requests (CreateTaskResult in SEP-1686).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[non_exhaustive]
pub struct CreateTaskResult {
    pub task: Task,
    #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
    pub meta: Option<Meta>,
}

impl CreateTaskResult {
    /// Create a new CreateTaskResult.
    pub fn new(task: Task) -> Self {
        Self { task, meta: None }
    }

    /// Sets the protocol-level metadata for this result.
    pub fn with_meta(mut self, meta: Meta) -> Self {
        self.meta = Some(meta);
        self
    }
}

/// Response to a `tasks/get` request.
///
/// Per spec, `GetTaskResult = allOf[Result, Task]` — the Task fields are
/// flattened at the top level, not nested under a `task` key.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[non_exhaustive]
pub struct GetTaskResult {
    #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")]
    pub meta: Option<Meta>,
    #[serde(flatten)]
    pub task: Task,
}

impl GetTaskResult {
    pub fn new(task: Task) -> Self {
        Self { meta: None, task }
    }
}

/// Response to a `tasks/result` request.
///
/// Per spec, the result structure matches the original request type
/// (e.g., `CallToolResult` for `tools/call`). This is represented as
/// an open object. The payload is the original request's result
/// serialized as a JSON value.
#[derive(Debug, Clone, PartialEq, Serialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[non_exhaustive]
pub struct GetTaskPayloadResult(pub Value);

impl GetTaskPayloadResult {
    /// Create a new GetTaskPayloadResult with the given value.
    pub fn new(value: Value) -> Self {
        Self(value)
    }
}

// Custom Deserialize that always fails, so that `GetTaskPayloadResult` is skipped
// during `#[serde(untagged)]` enum deserialization (e.g. `ServerResult`).
// The payload has the same JSON shape as `CustomResult(Value)`, so they are
// indistinguishable.  `CustomResult` acts as the catch-all instead.
// `GetTaskPayloadResult` should be constructed programmatically via `::new()`.
impl<'de> serde::Deserialize<'de> for GetTaskPayloadResult {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        // Consume the value so the deserializer state stays consistent.
        serde::de::IgnoredAny::deserialize(deserializer)?;
        Err(serde::de::Error::custom(
            "GetTaskPayloadResult cannot be deserialized directly; \
             use CustomResult as the catch-all",
        ))
    }
}

/// Response to a `tasks/cancel` request.
///
/// Per spec, `CancelTaskResult = allOf[Result, Task]` — same shape as `GetTaskResult`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[non_exhaustive]
pub struct CancelTaskResult {
    #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")]
    pub meta: Option<Meta>,
    #[serde(flatten)]
    pub task: Task,
}

impl CancelTaskResult {
    pub fn new(task: Task) -> Self {
        Self { meta: None, task }
    }
}