rmcp 3.0.0-beta.1

//! Task types for the MCP Tasks extension (SEP-2663).
//!
//! Tasks are defined by the official `io.modelcontextprotocol/tasks` extension.
//! A server may respond to a supported request (currently `tools/call`) with a
//! [`CreateTaskResult`] (`resultType: "task"`) instead of the standard result.
//! The client then polls `tasks/get`, answers in-task server-to-client requests
//! via `tasks/update`, and may signal cancellation via `tasks/cancel`.

use serde::{Deserialize, Serialize};

use super::{InputRequests, JsonObject, MetaObject, ResultType};

/// Extension identifier for the MCP Tasks extension (SEP-2663).
pub const TASKS_EXTENSION_ID: &str = "io.modelcontextprotocol/tasks";

/// Canonical task lifecycle status (SEP-2663).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[non_exhaustive]
pub enum TaskStatus {
    /// The request is currently being processed.
    #[default]
    Working,
    /// The server needs input from the client before the task can proceed.
    InputRequired,
    /// The request completed successfully and the result is available.
    /// This includes tool calls that returned results with `isError: true`.
    Completed,
    /// The request failed due to a JSON-RPC error during execution.
    Failed,
    /// The request was cancelled before completion.
    Cancelled,
}

impl TaskStatus {
    /// Returns `true` for terminal statuses (`completed`, `failed`, `cancelled`).
    pub fn is_terminal(&self) -> bool {
        matches!(self, Self::Completed | Self::Failed | Self::Cancelled)
    }
}

/// Operational metadata about ongoing work (spec `Task`, SEP-2663).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[non_exhaustive]
pub struct Task {
    /// Stable identifier for this task, generated by the server.
    pub task_id: String,
    /// Current task status.
    pub status: TaskStatus,
    /// Optional message describing the current task state.
    /// This MAY be exposed to the end-user or model.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub status_message: Option<String>,
    /// ISO 8601 timestamp when the task was created.
    pub created_at: String,
    /// ISO 8601 timestamp when the task was last updated.
    pub last_updated_at: String,
    /// Time-to-live duration from creation in integer milliseconds; `None`
    /// (serialized as `null`) means unlimited. The server may discard the task
    /// after the TTL elapses. This value MAY change over the lifetime of a task.
    pub ttl_ms: Option<u64>,
    /// Suggested polling interval in integer milliseconds. Clients SHOULD honor
    /// this value to avoid overwhelming the server. This value MAY change over
    /// the lifetime of a task.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub poll_interval_ms: Option<u64>,
}

impl Task {
    /// Create a new task with required fields.
    pub fn new(
        task_id: impl Into<String>,
        status: TaskStatus,
        created_at: impl Into<String>,
        last_updated_at: impl Into<String>,
    ) -> Self {
        Self {
            task_id: task_id.into(),
            status,
            status_message: None,
            created_at: created_at.into(),
            last_updated_at: last_updated_at.into(),
            ttl_ms: None,
            poll_interval_ms: 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_ms(mut self, ttl_ms: u64) -> Self {
        self.ttl_ms = Some(ttl_ms);
        self
    }

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

/// Status-specific payload carried alongside the base [`Task`] fields in a
/// [`DetailedTask`].
///
/// Mirrors the spec's `WorkingTask` / `InputRequiredTask` / `CompletedTask` /
/// `FailedTask` / `CancelledTask` union: the variant is discriminated by the
/// `status` field on the wire, with the payload fields inlined at the top level.
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
pub enum TaskPayload {
    /// `status: "working"` — no additional payload.
    Working,
    /// `status: "input_required"` — outstanding server-to-client requests.
    InputRequired {
        /// Server-to-client requests that need to be fulfilled during task
        /// execution. Keys are arbitrary identifiers for matching requests
        /// to responses, unique over the lifetime of the task.
        input_requests: InputRequests,
    },
    /// `status: "completed"` — the final result of the task. The structure
    /// matches the result type of the original request (e.g. `CallToolResult`).
    Completed {
        /// The final result of the original request.
        result: JsonObject,
    },
    /// `status: "failed"` — the JSON-RPC error that caused the task to fail.
    Failed {
        /// The JSON-RPC error object.
        error: JsonObject,
    },
    /// `status: "cancelled"` — no additional payload.
    Cancelled,
}

impl TaskPayload {
    /// The [`TaskStatus`] this payload corresponds to.
    pub fn status(&self) -> TaskStatus {
        match self {
            Self::Working => TaskStatus::Working,
            Self::InputRequired { .. } => TaskStatus::InputRequired,
            Self::Completed { .. } => TaskStatus::Completed,
            Self::Failed { .. } => TaskStatus::Failed,
            Self::Cancelled => TaskStatus::Cancelled,
        }
    }
}

/// A task with its status-specific payload inlined (spec `DetailedTask`).
///
/// Used by `tasks/get` responses ([`GetTaskResult`]) and `notifications/tasks`
/// ([`TaskStatusNotificationParams`](crate::model::TaskStatusNotificationParams)).
/// On the wire, the payload fields (`inputRequests` / `result` / `error`) are
/// flattened at the top level next to the base [`Task`] fields, and `status`
/// discriminates the variant.
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
pub struct DetailedTask {
    /// Base task metadata. Its `status` always agrees with the payload.
    pub task: Task,
    /// Status-specific payload.
    pub payload: TaskPayload,
}

impl DetailedTask {
    /// Build a `DetailedTask`, forcing `task.status` to match the payload.
    pub fn new(mut task: Task, payload: TaskPayload) -> Self {
        task.status = payload.status();
        Self { task, payload }
    }

    /// The current status.
    pub fn status(&self) -> TaskStatus {
        self.task.status
    }
}

// Wire shape helper: base Task fields + optional payload fields, all flattened.
#[derive(Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
struct DetailedTaskWire {
    #[serde(flatten)]
    task: Task,
    #[serde(skip_serializing_if = "Option::is_none")]
    input_requests: Option<InputRequests>,
    #[serde(skip_serializing_if = "Option::is_none")]
    result: Option<JsonObject>,
    #[serde(skip_serializing_if = "Option::is_none")]
    error: Option<JsonObject>,
}

impl From<DetailedTask> for DetailedTaskWire {
    fn from(value: DetailedTask) -> Self {
        let DetailedTask { task, payload } = value;
        let (input_requests, result, error) = match payload {
            TaskPayload::Working | TaskPayload::Cancelled => (None, None, None),
            TaskPayload::InputRequired { input_requests } => (Some(input_requests), None, None),
            TaskPayload::Completed { result } => (None, Some(result), None),
            TaskPayload::Failed { error } => (None, None, Some(error)),
        };
        Self {
            task,
            input_requests,
            result,
            error,
        }
    }
}

impl TryFrom<DetailedTaskWire> for DetailedTask {
    type Error = String;
    fn try_from(wire: DetailedTaskWire) -> Result<Self, String> {
        let payload = match wire.task.status {
            TaskStatus::Working => TaskPayload::Working,
            TaskStatus::Cancelled => TaskPayload::Cancelled,
            TaskStatus::InputRequired => TaskPayload::InputRequired {
                input_requests: wire.input_requests.ok_or_else(|| {
                    "task with status \"input_required\" is missing `inputRequests`".to_owned()
                })?,
            },
            TaskStatus::Completed => TaskPayload::Completed {
                result: wire.result.ok_or_else(|| {
                    "task with status \"completed\" is missing `result`".to_owned()
                })?,
            },
            TaskStatus::Failed => TaskPayload::Failed {
                error: wire
                    .error
                    .ok_or_else(|| "task with status \"failed\" is missing `error`".to_owned())?,
            },
        };
        Ok(DetailedTask {
            task: wire.task,
            payload,
        })
    }
}

impl Serialize for DetailedTask {
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        DetailedTaskWire::from(self.clone()).serialize(serializer)
    }
}

impl<'de> Deserialize<'de> for DetailedTask {
    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        let wire = DetailedTaskWire::deserialize(deserializer)?;
        Self::try_from(wire).map_err(serde::de::Error::custom)
    }
}

#[cfg(feature = "schemars")]
impl schemars::JsonSchema for DetailedTask {
    fn schema_name() -> std::borrow::Cow<'static, str> {
        "DetailedTask".into()
    }
    fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema {
        // The actual wire shape: base Task fields plus the optional
        // status-specific payload fields (inputRequests / result / error).
        <DetailedTaskWire as schemars::JsonSchema>::json_schema(generator)
    }
}

/// Result returned in lieu of a standard result to indicate the request will
/// be processed asynchronously (spec `CreateTaskResult`, `resultType: "task"`).
///
/// The embedded task is the seed state for the task; the client uses
/// `task.task_id` for all subsequent `tasks/get`, `tasks/update`, and
/// `tasks/cancel` calls.
#[derive(Debug, Clone, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[non_exhaustive]
pub struct CreateTaskResult {
    /// Always `"task"`.
    pub result_type: ResultType,
    /// Seed state of the newly created task, flattened at the top level.
    #[serde(flatten)]
    pub task: Task,
    #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")]
    pub meta: Option<MetaObject>,
}

// Custom deserializer that requires `resultType: "task"`. Without this,
// `CreateTaskResult` would greedily match other task-shaped results (e.g.
// `tasks/get` responses, which also carry `taskId`/`status` at the top level
// but use `resultType: "complete"`) inside `#[serde(untagged)]` unions such
// as `ServerResult`.
impl<'de> Deserialize<'de> for CreateTaskResult {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        #[derive(Deserialize)]
        #[serde(rename_all = "camelCase")]
        struct Helper {
            result_type: ResultType,
            #[serde(flatten)]
            task: Task,
            #[serde(rename = "_meta", default)]
            meta: Option<MetaObject>,
        }
        let helper = Helper::deserialize(deserializer)?;
        if !helper.result_type.is_task() {
            return Err(serde::de::Error::custom(
                "CreateTaskResult requires resultType to be \"task\"",
            ));
        }
        Ok(CreateTaskResult {
            result_type: helper.result_type,
            task: helper.task,
            meta: helper.meta,
        })
    }
}

impl CreateTaskResult {
    /// Create a new `CreateTaskResult` from the seed task state.
    pub fn new(task: Task) -> Self {
        Self {
            result_type: ResultType::TASK,
            task,
            meta: None,
        }
    }

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

/// Response to a `tasks/get` request (spec `GetTaskResult = Result & DetailedTask`).
///
/// `resultType` is `"complete"` — this is the standard result shape for
/// `tasks/get`, not a task handle.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[non_exhaustive]
pub struct GetTaskResult {
    /// Result type discriminator. `tasks/get` responses are standard results:
    /// `"complete"` (SEP-2322). Absent values deserialize as `"complete"`.
    #[serde(default)]
    pub result_type: ResultType,
    #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")]
    pub meta: Option<MetaObject>,
    /// The task with status-specific payload inlined.
    #[serde(flatten)]
    pub task: DetailedTask,
}

impl GetTaskResult {
    pub fn new(task: DetailedTask) -> Self {
        Self {
            result_type: ResultType::COMPLETE,
            meta: None,
            task,
        }
    }
}

/// Empty acknowledgement for `tasks/update` and `tasks/cancel` (SEP-2663).
///
/// The spec requires these acks to be empty results carrying the SEP-2322
/// `resultType: "complete"` discriminator; task state changes are observed
/// via the next `tasks/get`.
#[derive(Debug, Clone, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[non_exhaustive]
pub struct TaskAckResult {
    /// Always `"complete"`.
    pub result_type: ResultType,
    #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")]
    pub meta: Option<MetaObject>,
}

// Custom deserializer that requires `resultType: "complete"` and rejects any
// other fields. Without this, `TaskAckResult` would greedily match arbitrary
// result objects carrying a `resultType` key inside `#[serde(untagged)]`
// unions such as `ServerResult`, shadowing `CustomResult` and losing data.
impl<'de> Deserialize<'de> for TaskAckResult {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        #[derive(Deserialize)]
        #[serde(rename_all = "camelCase", deny_unknown_fields)]
        struct Helper {
            result_type: ResultType,
            #[serde(rename = "_meta", default)]
            meta: Option<MetaObject>,
        }
        let helper = Helper::deserialize(deserializer)?;
        if !helper.result_type.is_complete() {
            return Err(serde::de::Error::custom(
                "TaskAckResult requires resultType to be \"complete\"",
            ));
        }
        Ok(TaskAckResult {
            result_type: helper.result_type,
            meta: helper.meta,
        })
    }
}

impl Default for TaskAckResult {
    fn default() -> Self {
        Self {
            result_type: ResultType::COMPLETE,
            meta: None,
        }
    }
}

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

#[cfg(test)]
mod tests {
    use serde_json::json;

    use super::*;

    fn base_task(status: TaskStatus) -> Task {
        Task::new(
            "task-1",
            status,
            "2025-11-25T10:30:00Z",
            "2025-11-25T10:40:00Z",
        )
        .with_ttl_ms(60000)
        .with_poll_interval_ms(5000)
    }

    #[test]
    fn create_task_result_wire_shape() {
        let result = CreateTaskResult::new(base_task(TaskStatus::Working));
        let value = serde_json::to_value(&result).unwrap();
        assert_eq!(
            value,
            json!({
                "resultType": "task",
                "taskId": "task-1",
                "status": "working",
                "createdAt": "2025-11-25T10:30:00Z",
                "lastUpdatedAt": "2025-11-25T10:40:00Z",
                "ttlMs": 60000,
                "pollIntervalMs": 5000
            })
        );
        let roundtrip: CreateTaskResult = serde_json::from_value(value).unwrap();
        assert_eq!(roundtrip, result);
    }

    #[test]
    fn ttl_ms_null_means_unlimited() {
        let mut task = base_task(TaskStatus::Working);
        task.ttl_ms = None;
        let value = serde_json::to_value(&task).unwrap();
        assert_eq!(value["ttlMs"], serde_json::Value::Null);
        let roundtrip: Task = serde_json::from_value(value).unwrap();
        assert_eq!(roundtrip.ttl_ms, None);
    }

    #[test]
    fn detailed_task_completed_roundtrip() {
        let detailed = DetailedTask::new(
            base_task(TaskStatus::Working),
            TaskPayload::Completed {
                result: serde_json::from_value(json!({
                    "content": [{"type": "text", "text": "ok"}],
                    "isError": false
                }))
                .unwrap(),
            },
        );
        // Status is forced to match the payload.
        assert_eq!(detailed.status(), TaskStatus::Completed);
        let value = serde_json::to_value(&detailed).unwrap();
        assert_eq!(value["status"], "completed");
        assert_eq!(value["result"]["isError"], false);
        let roundtrip: DetailedTask = serde_json::from_value(value).unwrap();
        assert_eq!(roundtrip, detailed);
    }

    #[test]
    fn detailed_task_input_required_requires_input_requests() {
        let err = serde_json::from_value::<DetailedTask>(json!({
            "taskId": "task-1",
            "status": "input_required",
            "createdAt": "2025-11-25T10:30:00Z",
            "lastUpdatedAt": "2025-11-25T10:40:00Z",
            "ttlMs": null
        }))
        .unwrap_err();
        assert!(err.to_string().contains("inputRequests"));
    }

    #[test]
    fn detailed_task_failed_roundtrip() {
        let detailed = DetailedTask::new(
            base_task(TaskStatus::Failed),
            TaskPayload::Failed {
                error: serde_json::from_value(json!({
                    "code": -32603,
                    "message": "boom"
                }))
                .unwrap(),
            },
        );
        let value = serde_json::to_value(&detailed).unwrap();
        assert_eq!(value["status"], "failed");
        assert_eq!(value["error"]["code"], -32603);
        let roundtrip: DetailedTask = serde_json::from_value(value).unwrap();
        assert_eq!(roundtrip, detailed);
    }

    #[test]
    fn get_task_result_flattens_detailed_task() {
        let result = GetTaskResult::new(DetailedTask::new(
            base_task(TaskStatus::Working),
            TaskPayload::Working,
        ));
        let value = serde_json::to_value(&result).unwrap();
        assert_eq!(value["taskId"], "task-1");
        assert_eq!(value["status"], "working");
        let roundtrip: GetTaskResult = serde_json::from_value(value).unwrap();
        assert_eq!(roundtrip, result);
    }
}