neva 0.5.0

//! Types and utilities for tracking tasks

use crate::error::{Error, ErrorCode};
use crate::types::{Task, TaskStatus};
use chrono::Utc;
use std::{
    cmp::Ordering,
    collections::BinaryHeap,
    sync::{
        Mutex,
        atomic::{AtomicU64, Ordering as AtomicOrdering},
    },
};
use tokio_util::sync::{CancellationToken, WaitForCancellationFuture};
#[cfg(feature = "legacy-spec")]
use {
    crate::types::TaskPayload,
    serde::Serialize,
    tokio::sync::watch::{Receiver, Sender, channel},
};

#[cfg(all(feature = "server", not(feature = "legacy-spec")))]
use tokio::sync::oneshot;

pub(crate) struct TaskTracker {
    tasks: dashmap::DashMap<String, TaskEntry>,
    expirations: Mutex<BinaryHeap<TaskExpiry>>,
    next_expiry_seq: AtomicU64,
}

/// Alias for [`Option<TaskPayload>`]
#[cfg(feature = "legacy-spec")]
pub(crate) type MaybePayload = Option<TaskPayload>;

/// Represents a task currently running on the server
pub(crate) struct TaskEntry {
    task: Task,
    token: CancellationToken,
    // The `tasks/result` watch channel. MCP 2026-07-28 has no `tasks/result`:
    // the outcome is stored on the entry and reported by `tasks/get`.
    #[cfg(all(feature = "server", feature = "legacy-spec"))]
    tx: Sender<MaybePayload>,
    #[cfg(feature = "legacy-spec")]
    rx: Receiver<MaybePayload>,
    /// MCP 2026-07-28 state that `tasks/get` reports beyond the bare status:
    /// the terminal outcome and any outstanding input requests.
    #[cfg(all(feature = "server", not(feature = "legacy-spec")))]
    state: Mutex<TaskState>,
}

/// The status-specific half of a [`DetailedTask`](crate::types::DetailedTask).
///
/// A task carries at most one of these at a time: pending `inputs` while it is
/// `input_required`, then either `result` or `error` once it terminates.
#[cfg(all(feature = "server", not(feature = "legacy-spec")))]
#[derive(Default)]
struct TaskState {
    /// Final result, set when the task completes successfully.
    result: Option<serde_json::Value>,

    /// JSON-RPC error, set when the task fails.
    error: Option<serde_json::Value>,

    /// Outstanding input requests, keyed as the client will echo them back in
    /// `tasks/update`.
    inputs: std::collections::HashMap<String, PendingInput>,
}

/// One outstanding input request plus the slot its answer resumes.
///
/// Under the stateless transport a suspended `ctx.task().elicit` cannot be
/// correlated by session (a fresh one is minted per POST), so the answer
/// arrives out-of-band as a `tasks/update` addressed to the task id and keyed
/// by the same key `tasks/get` surfaced.
#[cfg(all(feature = "server", not(feature = "legacy-spec")))]
struct PendingInput {
    request: crate::types::mrtr::InputRequest,
    tx: oneshot::Sender<serde_json::Value>,
}

/// Represents a handle to a task that can be used to cancel or get the result of the task.
pub(crate) struct TaskHandle {
    token: CancellationToken,
    #[cfg(feature = "legacy-spec")]
    tx: Sender<MaybePayload>,
}

struct TaskExpiry {
    deadline_ms: i64,
    sequence: u64,
    id: String,
}

impl PartialEq for TaskExpiry {
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        self.deadline_ms == other.deadline_ms && self.sequence == other.sequence
    }
}

impl Eq for TaskExpiry {}

impl PartialOrd for TaskExpiry {
    #[inline]
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for TaskExpiry {
    #[inline]
    fn cmp(&self, other: &Self) -> Ordering {
        other
            .deadline_ms
            .cmp(&self.deadline_ms)
            .then_with(|| other.sequence.cmp(&self.sequence))
    }
}

impl TaskTracker {
    /// Creates a new [`TaskTracker`]
    #[inline]
    pub(crate) fn new() -> Self {
        Self {
            tasks: dashmap::DashMap::new(),
            expirations: Mutex::new(BinaryHeap::new()),
            next_expiry_seq: AtomicU64::new(0),
        }
    }

    /// Returns a list of currently running tasks.
    #[cfg(feature = "legacy-spec")]
    pub(crate) fn tasks(&self) -> Vec<Task> {
        self.cleanup_expired();

        self.tasks
            .iter()
            .map(|entry| entry.task.clone())
            .collect::<Vec<_>>()
    }

    /// Tacks the task and returns the [`TaskHandle`] for this task
    pub(crate) fn track(&self, task: Task) -> TaskHandle {
        self.cleanup_expired();

        let token = CancellationToken::new();
        #[cfg(feature = "legacy-spec")]
        let (tx, rx) = channel(None);

        self.tasks.insert(
            task.id.clone(),
            TaskEntry {
                token: token.clone(),
                #[cfg(all(feature = "server", feature = "legacy-spec"))]
                tx: tx.clone(),
                task,
                #[cfg(feature = "legacy-spec")]
                rx,
                #[cfg(all(feature = "server", not(feature = "legacy-spec")))]
                state: Mutex::new(TaskState::default()),
            },
        );

        TaskHandle {
            token,
            #[cfg(feature = "legacy-spec")]
            tx,
        }
    }

    /// Cancels the task
    ///
    /// Cancellation is cooperative under MCP 2026-07-28: the entry stays until
    /// its TTL expires so the requestor can keep polling `tasks/get` for the
    /// outcome. A tool that wins the race with the cancellation token still
    /// overwrites the status with its own terminal one, which the spec allows.
    #[cfg(not(feature = "legacy-spec"))]
    pub(crate) fn cancel(&self, id: &str) -> Result<Task, Error> {
        self.cleanup_expired();

        if let Some(mut entry) = self.tasks.get_mut(id) {
            // A terminal task is not cancellable -- cancelling asks a task to
            // stop, it does not rewrite one that already stopped. Overwriting
            // the status here would also hide the outcome, since `get_state`
            // reports a `result` or `error` only under the status that owns it.
            // Retention makes this reachable in two ways: a cancel racing the
            // tool's own completion, and one sent against a task that finished
            // long ago but is still inside its TTL.
            if Self::is_terminal(&entry.task) {
                return Ok(entry.task.clone());
            }
            entry.token.cancel();
            entry.task = entry.task.clone().cancel();
            self.schedule_expiry(&entry.task);
            Ok(entry.task.clone())
        } else {
            Err(Error::new(
                ErrorCode::InvalidParams,
                format!("Could not find task with id: {id}"),
            ))
        }
    }

    /// Cancels the task
    #[cfg(feature = "legacy-spec")]
    pub(crate) fn cancel(&self, id: &str) -> Result<Task, Error> {
        self.cleanup_expired();

        if let Some((_, entry)) = self.tasks.remove(id) {
            entry.token.cancel();
            Ok(entry.task.cancel())
        } else {
            Err(Error::new(
                ErrorCode::InvalidParams,
                format!("Could not find task with id: {id}"),
            ))
        }
    }

    /// Completes the task
    pub(crate) fn complete(&self, id: &str) {
        self.cleanup_expired();

        if let Some(mut entry) = self.tasks.get_mut(id) {
            entry.task.complete();
            self.schedule_expiry(&entry.task);
        }
    }

    /// Fails the task
    #[cfg(feature = "server")]
    pub(crate) fn fail(&self, id: &str) {
        self.cleanup_expired();

        if let Some(mut entry) = self.tasks.get_mut(id) {
            entry.task.fail();
            self.schedule_expiry(&entry.task);
        }
    }

    /// Sets the task into `input_required` status
    #[cfg(feature = "server")]
    #[cfg_attr(not(feature = "legacy-spec"), allow(dead_code))]
    pub(crate) fn require_input(&self, id: &str) {
        self.cleanup_expired();

        if let Some(mut entry) = self.tasks.get_mut(id) {
            entry.task.require_input();
        }
    }

    /// Sets the task into `working` status
    #[cfg(all(feature = "server", feature = "legacy-spec"))]
    pub(crate) fn reset(&self, id: &str) {
        self.cleanup_expired();

        if let Some(mut entry) = self.tasks.get_mut(id) {
            entry.task.reset();
            let _ = entry.tx.send(None);
        }
    }

    /// Sets the result of the [`Task`].
    #[cfg(all(feature = "server", feature = "legacy-spec"))]
    pub(crate) fn set_result<T: Serialize>(&self, id: &str, result: T) {
        self.cleanup_expired();

        if let Some(entry) = self.tasks.get(id) {
            let result = match serde_json::to_value(result) {
                Ok(result) => result,
                Err(_err) => {
                    #[cfg(feature = "tracing")]
                    tracing::error!(logger = "neva", "Unable to serialize task result: {_err:?}");
                    return;
                }
            };
            let _ = entry.tx.send(Some(TaskPayload(result)));
        }
    }

    /// Records an outstanding input request under `key` and returns the
    /// receiver a later [`Self::provide_inputs`] fulfills.
    ///
    /// Replaces any previously parked (unanswered) request under the same key.
    /// Returns `None` if the task no longer exists.
    #[cfg(all(feature = "server", not(feature = "legacy-spec")))]
    pub(crate) fn park_input(
        &self,
        id: &str,
        key: String,
        request: crate::types::mrtr::InputRequest,
    ) -> Option<oneshot::Receiver<serde_json::Value>> {
        let entry = self.tasks.get(id)?;
        let (tx, rx) = oneshot::channel();
        match entry.state.lock() {
            Ok(mut state) => {
                state.inputs.insert(key, PendingInput { request, tx });
                Some(rx)
            }
            Err(_) => None,
        }
    }

    /// Delivers client answers to a task's outstanding input requests
    /// (`tasks/update`).
    ///
    /// Keys that match nothing outstanding are ignored, per the spec -- a
    /// retried `tasks/update` must not fail. The task returns to `working` once
    /// no requests are left outstanding.
    #[cfg(all(feature = "server", not(feature = "legacy-spec")))]
    pub(crate) fn provide_inputs(
        &self,
        id: &str,
        responses: crate::types::mrtr::InputResponses,
    ) -> Result<(), Error> {
        self.cleanup_expired();

        let mut entry = self.tasks.get_mut(id).ok_or_else(|| {
            Error::new(
                ErrorCode::InvalidParams,
                format!("Could not find task with id: {id}"),
            )
        })?;

        let outstanding = {
            let Ok(mut state) = entry.state.lock() else {
                return Err(Error::new(
                    ErrorCode::InternalError,
                    "Unable to update task inputs",
                ));
            };
            for (key, value) in responses {
                if let Some(pending) = state.inputs.remove(&key) {
                    let _ = pending.tx.send(value);
                }
            }
            state.inputs.len()
        };

        // Still waiting on other keys: leave the task in `input_required`.
        if outstanding == 0 && entry.task.status == TaskStatus::InputRequired {
            entry.task.reset();
        }
        Ok(())
    }

    /// Stores the terminal outcome reported by `tasks/get`.
    ///
    /// `Ok` lands in `result` (status `completed`), `Err` in `error`
    /// (status `failed`) as a JSON-RPC error object.
    ///
    /// The stored result *is* the result the original request would have
    /// returned synchronously, so it carries the same mandatory `resultType`
    /// discriminator -- it just never passes through
    /// [`Response::success`](crate::types::Response::success), which is where
    /// a directly-returned result is stamped.
    #[cfg(all(feature = "server", not(feature = "legacy-spec")))]
    pub(crate) fn set_outcome(&self, id: &str, outcome: Result<serde_json::Value, Error>) {
        self.cleanup_expired();

        let Some(entry) = self.tasks.get(id) else {
            return;
        };
        let Ok(mut state) = entry.state.lock() else {
            return;
        };
        match outcome {
            Ok(result) => state.result = Some(crate::types::tag_complete(result)),
            Err(err) => {
                state.error = serde_json::to_value(crate::types::ErrorDetails::from(err)).ok()
            }
        }
    }

    /// Retrieves the full task state served by `tasks/get`: the status plus the
    /// outstanding input requests, the terminal result, or the error.
    #[cfg(all(feature = "server", not(feature = "legacy-spec")))]
    pub(crate) fn get_state(&self, id: &str) -> Result<crate::types::DetailedTask, Error> {
        self.cleanup_expired();

        let entry = self.tasks.get(id).ok_or_else(|| {
            Error::new(
                ErrorCode::InvalidParams,
                format!("Could not find task with id: {id}"),
            )
        })?;

        let mut detailed = crate::types::DetailedTask::from(entry.task.clone());
        if let Ok(state) = entry.state.lock() {
            // The wire shape is discriminated by status: `inputRequests` belongs
            // to `input_required`, `result` to `completed`, `error` to `failed`,
            // and a cancelled task carries none of them. The state can outlive
            // the status it was recorded under -- a task cancelled while parked
            // on `ctx.task().elicit` keeps its pending input, and a tool that
            // races the cancellation records an outcome -- so what is reported
            // is decided by the status, not by what happens to be stored.
            match detailed.status {
                TaskStatus::InputRequired if !state.inputs.is_empty() => {
                    detailed.input_requests = Some(
                        state
                            .inputs
                            .iter()
                            .map(|(key, pending)| (key.clone(), pending.request.clone()))
                            .collect(),
                    );
                }
                TaskStatus::Completed => detailed.result = state.result.clone(),
                TaskStatus::Failed => detailed.error = state.error.clone(),
                _ => {}
            }
        }
        Ok(detailed)
    }

    /// Retrieves the task status
    ///
    /// The status alone is a legacy shape: MCP 2026-07-28 reports a task -- to
    /// `tasks/get` and to `notifications/tasks` alike -- as the full
    /// [`DetailedTask`](crate::types::DetailedTask) that [`Self::get_state`]
    /// builds.
    #[cfg(any(feature = "legacy-spec", test))]
    pub(crate) fn get_status(&self, id: &str) -> Result<Task, Error> {
        self.cleanup_expired();

        self.tasks.get(id).map(|t| t.task.clone()).ok_or_else(|| {
            Error::new(
                ErrorCode::InvalidParams,
                format!("Could not find task with id: {id}"),
            )
        })
    }

    /// Returns the task result if it is present,
    /// otherwise waits until the result is available or the task will be canceled.
    #[cfg(feature = "legacy-spec")]
    pub(crate) async fn get_result(&self, id: &str) -> Result<TaskPayload, Error> {
        self.cleanup_expired();

        let (status, mut result_rx, token) = {
            let entry = self.tasks.get(id).ok_or_else(|| {
                Error::new(
                    ErrorCode::InvalidParams,
                    format!("Could not find task with id: {id}"),
                )
            })?;

            (entry.task.status, entry.rx.clone(), entry.token.clone())
        };

        if let Some(ref result) = *result_rx.borrow_and_update() {
            if status != TaskStatus::InputRequired {
                self.tasks.remove(id);
            }
            return Ok(result.clone());
        }

        loop {
            tokio::select! {
                changed = result_rx.changed() => {
                    if changed.is_err() {
                        return Err(Error::new(ErrorCode::InternalError, "Unable to get task result"));
                    }

                    if let Some(result) = result_rx.borrow_and_update().clone() {
                        let task = self.get_status(id)?;
                        if task.status != TaskStatus::InputRequired {
                            self.tasks.remove(id);
                        }
                        return Ok(result);
                    }
                }
                _ = token.cancelled() => {
                    return Err(Error::new(ErrorCode::InvalidRequest, "Task has been cancelled"));
                }
            }
        }
    }

    #[inline]
    fn cleanup_expired(&self) {
        let now_ms = Utc::now().timestamp_millis();
        let mut expired = Vec::new();

        if let Ok(mut expirations) = self.expirations.lock() {
            while expirations
                .peek()
                .is_some_and(|entry| entry.deadline_ms <= now_ms)
            {
                let entry = expirations.pop().expect("peeked entry must exist");
                expired.push((entry.id, entry.deadline_ms));
            }
        }

        for (id, deadline_ms) in expired {
            let should_remove = self.tasks.get(&id).is_some_and(|entry| {
                Self::is_terminal(&entry.task)
                    && Self::task_deadline_ms(&entry.task).is_some_and(|d| d == deadline_ms)
            });

            if should_remove {
                let _ = self.tasks.remove(&id);
            }
        }
    }

    #[inline]
    fn is_terminal(task: &Task) -> bool {
        matches!(
            task.status,
            TaskStatus::Completed | TaskStatus::Failed | TaskStatus::Cancelled
        )
    }

    #[inline]
    fn schedule_expiry(&self, task: &Task) {
        let Some(deadline_ms) = Self::task_deadline_ms(task) else {
            return;
        };

        let sequence = self.next_expiry_seq.fetch_add(1, AtomicOrdering::Relaxed);
        if let Ok(mut expirations) = self.expirations.lock() {
            expirations.push(TaskExpiry {
                deadline_ms,
                sequence,
                id: task.id.clone(),
            });
        }
    }

    #[inline]
    fn task_deadline_ms(task: &Task) -> Option<i64> {
        // `None` means unlimited retention: nothing to schedule.
        let ttl_ms = i64::try_from(task.ttl?).unwrap_or(i64::MAX);
        task.created_at.timestamp_millis().checked_add(ttl_ms)
    }
}

impl Default for TaskTracker {
    #[inline]
    fn default() -> Self {
        Self::new()
    }
}

impl TaskHandle {
    /// Completes the [`Task`] and sets the result.
    #[cfg(feature = "legacy-spec")]
    pub(crate) fn set_result<T: Serialize>(self, result: T) {
        let result = match serde_json::to_value(result) {
            Ok(result) => result,
            Err(_err) => {
                #[cfg(feature = "tracing")]
                tracing::error!(logger = "neva", "Unable to serialize task result: {_err:?}");
                return;
            }
        };
        let _ = self.tx.send(Some(TaskPayload(result)));
    }

    /// Returns a [`Future`] that gets fulfilled when cancellation is requested.
    ///
    /// Equivalent to:
    ///
    /// ```ignore
    /// async fn cancelled(&self);
    /// ```
    ///
    /// The future will complete immediately if the token is already canceled
    /// when this method is called.
    ///
    /// # Cancellation safety
    ///
    /// This method is cancel safe.
    #[inline]
    pub(crate) fn cancelled(&self) -> WaitForCancellationFuture<'_> {
        self.token.cancelled()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::TaskStatus;
    use std::sync::Arc;

    #[cfg(all(feature = "server", feature = "legacy-spec"))]
    use crate::types::CallToolResponse;

    #[cfg(feature = "legacy-spec")]
    #[test]
    fn it_can_create_new_tracker() {
        let tracker = TaskTracker::new();
        assert_eq!(tracker.tasks().len(), 0);
    }

    #[test]
    fn it_can_track_task() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();

        let _handle = tracker.track(task);

        assert_eq!(tracker.get_status(&task_id).unwrap().id, task_id);
    }

    #[cfg(feature = "legacy-spec")]
    #[test]
    fn it_can_return_list_of_tasks() {
        let tracker = TaskTracker::new();
        let task1 = Task::new();
        let task2 = Task::new();

        let _handle1 = tracker.track(task1.clone());
        let _handle2 = tracker.track(task2.clone());

        let tasks = tracker.tasks();
        assert_eq!(tasks.len(), 2);
    }

    #[cfg(feature = "legacy-spec")]
    #[test]
    fn it_can_cancel_task() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();

        let _handle = tracker.track(task);

        let result = tracker.cancel(&task_id).unwrap();
        assert_eq!(result.status, TaskStatus::Cancelled);
        assert!(tracker.get_status(&task_id).is_err());
    }

    #[cfg(not(feature = "legacy-spec"))]
    #[test]
    fn it_can_cancel_task() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();

        let _handle = tracker.track(task);

        let result = tracker.cancel(&task_id).unwrap();
        assert_eq!(result.status, TaskStatus::Cancelled);

        // The entry survives cancellation: the requestor learns the outcome by
        // polling `tasks/get`, which would otherwise answer `InvalidParams`.
        let status = tracker.get_status(&task_id).expect("cancelled task kept");
        assert_eq!(status.status, TaskStatus::Cancelled);
    }

    #[cfg(not(feature = "legacy-spec"))]
    #[test]
    fn it_does_remove_expired_cancelled_tasks() {
        let tracker = TaskTracker::new();
        let task = Task::from(crate::types::TaskMetadata { ttl: Some(1) });
        let task_id = task.id.clone();

        let _handle = tracker.track(task);
        tracker.cancel(&task_id).unwrap();

        std::thread::sleep(std::time::Duration::from_millis(10));

        assert!(tracker.get_status(&task_id).is_err());
    }

    #[test]
    fn it_does_return_error_when_cancelling_nonexistent_task() {
        let tracker = TaskTracker::new();

        let result = tracker.cancel("nonexistent");
        assert!(result.is_err());
        assert_eq!(result.unwrap_err().code, ErrorCode::InvalidParams);
    }

    #[test]
    fn it_can_complete_task() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();

        let _handle = tracker.track(task);

        tracker.complete(&task_id);

        let status = tracker.get_status(&task_id).unwrap();
        assert_eq!(status.status, TaskStatus::Completed);
    }

    #[test]
    fn it_does_nothing_when_completing_nonexistent_task() {
        let tracker = TaskTracker::new();
        tracker.complete("nonexistent");
        // Should not panic
    }

    #[cfg(feature = "server")]
    #[test]
    fn it_can_fail_task() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();

        let _handle = tracker.track(task);

        tracker.fail(&task_id);

        let status = tracker.get_status(&task_id).unwrap();
        assert_eq!(status.status, TaskStatus::Failed);
    }

    #[cfg(feature = "server")]
    #[test]
    fn it_does_nothing_when_failing_nonexistent_task() {
        let tracker = TaskTracker::new();
        tracker.fail("nonexistent");
        // Should not panic
    }

    #[cfg(feature = "server")]
    #[test]
    fn it_can_require_input() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();

        let _handle = tracker.track(task);

        tracker.require_input(&task_id);

        let status = tracker.get_status(&task_id).unwrap();
        assert_eq!(status.status, TaskStatus::InputRequired);
    }

    #[cfg(feature = "server")]
    #[test]
    fn it_does_nothing_when_requiring_input_for_nonexistent_task() {
        let tracker = TaskTracker::new();
        tracker.require_input("nonexistent");
        // Should not panic
    }

    #[test]
    fn it_can_get_task_status() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();

        let _handle = tracker.track(task.clone());

        let status = tracker.get_status(&task_id).unwrap();
        assert_eq!(status.id, task.id);
        assert_eq!(status.status, TaskStatus::Working);
    }

    #[test]
    fn it_does_return_error_when_getting_status_of_nonexistent_task() {
        let tracker = TaskTracker::new();

        let result = tracker.get_status("nonexistent");
        assert!(result.is_err());
        assert_eq!(result.unwrap_err().code, ErrorCode::InvalidParams);
    }

    #[cfg(feature = "legacy-spec")]
    #[tokio::test]
    async fn it_can_get_task_result_when_completed() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();

        let handle = tracker.track(task.clone());

        tokio::spawn(async move {
            tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
            handle.set_result("test_result".to_string());
        });

        let result = tracker.get_result(&task_id).await.unwrap();
        assert_eq!(result.0, "test_result");
    }

    #[cfg(feature = "legacy-spec")]
    #[tokio::test]
    async fn it_does_return_result_immediately_when_already_available() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();

        let handle = tracker.track(task.clone());
        handle.set_result("immediate_result".to_string());

        let result = tracker.get_result(&task_id).await.unwrap();
        assert_eq!(result.0, "immediate_result");
    }

    #[cfg(feature = "legacy-spec")]
    #[tokio::test]
    async fn it_does_return_error_when_getting_result_of_nonexistent_task() {
        let tracker = TaskTracker::new();

        let result = tracker.get_result("nonexistent").await;
        assert!(result.is_err());
        assert_eq!(result.unwrap_err().code, ErrorCode::InvalidParams);
    }

    #[cfg(feature = "legacy-spec")]
    #[tokio::test]
    async fn it_does_return_error_when_task_is_cancelled() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();

        let _handle = tracker.track(task.clone());

        let tracker = Arc::new(tracker);

        tokio::spawn({
            let tracker = tracker.clone();
            let task_id = task_id.clone();
            async move {
                tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
                let _ = tracker.cancel(&task_id);
            }
        });

        let result = tracker.get_result(&task_id).await;
        assert!(result.is_err());
        assert_eq!(result.unwrap_err().code, ErrorCode::InvalidRequest);
    }

    #[cfg(feature = "legacy-spec")]
    #[tokio::test]
    async fn it_can_wait_for_result_with_multiple_updates() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();

        let handle = tracker.track(task.clone());

        tokio::spawn(async move {
            tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
            handle.set_result("final_result".to_string());
        });

        let result = tracker.get_result(&task_id).await.unwrap();
        assert_eq!(result.0, "final_result");
        assert_eq!(tracker.tasks().len(), 0);
    }

    #[cfg(feature = "legacy-spec")]
    #[tokio::test]
    async fn it_does_remove_task_after_getting_result() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();

        let handle = tracker.track(task.clone());
        handle.set_result("result".to_string());

        let _ = tracker.get_result(&task_id).await.unwrap();

        assert_eq!(tracker.tasks().len(), 0);
    }

    #[tokio::test]
    async fn it_can_create_task_handle() {
        let tracker = TaskTracker::new();
        let task = Task::new();

        let handle = tracker.track(task);

        // Just ensure the handle can be created and used
        tokio::spawn(async move {
            tokio::select! {
                _ = handle.cancelled() => {}
            }
        });
    }

    #[tokio::test]
    async fn it_can_cancel_via_handle() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();

        let handle = tracker.track(task.clone());

        let tracker = Arc::new(tracker);

        tokio::spawn({
            let tracker = tracker.clone();
            let task_id = task_id.clone();
            async move {
                tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
                let _ = tracker.cancel(&task_id);
            }
        });

        tokio::select! {
            _ = handle.cancelled() => {
                // Successfully cancelled
            }
            _ = tokio::time::sleep(tokio::time::Duration::from_secs(1)) => {
                panic!("Task was not cancelled");
            }
        }
    }

    #[cfg(feature = "legacy-spec")]
    #[test]
    #[cfg(feature = "server")]
    fn it_can_handle_complex_payload_types() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();

        let handle = tracker.track(task.clone());

        let response = CallToolResponse::new("test");
        tracker.complete(&task_id);
        handle.set_result(response.clone());

        let status = tracker.get_status(&task_id).unwrap();
        assert_eq!(status.status, TaskStatus::Completed);
    }

    #[cfg(feature = "legacy-spec")]
    #[tokio::test]
    async fn it_can_track_multiple_concurrent_tasks() {
        let tracker = TaskTracker::new();
        let tasks: Vec<_> = (0..5).map(|_| Task::new()).collect();
        let task_ids: Vec<_> = tasks.iter().map(|t| t.id.clone()).collect();

        let handles: Vec<_> = tasks.into_iter().map(|t| tracker.track(t)).collect();

        for (i, handle) in handles.into_iter().enumerate() {
            let result = format!("result_{}", i);
            handle.set_result(result);
        }

        for (i, task_id) in task_ids.iter().enumerate() {
            let result = tracker.get_result(task_id).await.unwrap();
            assert_eq!(result.0, format!("result_{}", i));
        }

        assert_eq!(tracker.tasks().len(), 0);
    }

    #[cfg(all(feature = "server", not(feature = "legacy-spec")))]
    #[tokio::test]
    async fn park_input_is_fulfilled_by_tasks_update() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();
        let _handle = tracker.track(task);

        let rx = tracker
            .park_input(&task_id, "k1".into(), elicit_request())
            .expect("parked");
        tracker.require_input(&task_id);

        // `tasks/get` surfaces the outstanding ask under the same key.
        let state = tracker.get_state(&task_id).unwrap();
        assert_eq!(state.status, TaskStatus::InputRequired);
        assert!(state.input_requests.expect("asks").contains_key("k1"));

        let mut answers = crate::types::mrtr::InputResponses::new();
        answers.insert("k1".into(), serde_json::json!({ "action": "accept" }));
        tracker.provide_inputs(&task_id, answers).unwrap();

        let got = rx.await.expect("answer delivered");
        assert_eq!(got, serde_json::json!({ "action": "accept" }));

        // No asks left outstanding -> back to working, nothing surfaced.
        let state = tracker.get_state(&task_id).unwrap();
        assert_eq!(state.status, TaskStatus::Working);
        assert!(state.input_requests.is_none());
    }

    #[cfg(all(feature = "server", not(feature = "legacy-spec")))]
    #[test]
    fn provide_inputs_ignores_unknown_keys() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();
        let _handle = tracker.track(task);

        let _rx = tracker
            .park_input(&task_id, "k1".into(), elicit_request())
            .expect("parked");
        tracker.require_input(&task_id);

        // A retried or stale `tasks/update` must not fail...
        let mut answers = crate::types::mrtr::InputResponses::new();
        answers.insert("nope".into(), serde_json::json!({}));
        assert!(tracker.provide_inputs(&task_id, answers).is_ok());

        // ...and must not resolve the still-outstanding ask.
        let state = tracker.get_state(&task_id).unwrap();
        assert_eq!(state.status, TaskStatus::InputRequired);
        assert!(state.input_requests.expect("asks").contains_key("k1"));
    }

    #[cfg(all(feature = "server", not(feature = "legacy-spec")))]
    #[test]
    fn provide_inputs_errors_for_unknown_task() {
        let tracker = TaskTracker::new();
        let result = tracker.provide_inputs("nonexistent", Default::default());

        assert_eq!(result.unwrap_err().code, ErrorCode::InvalidParams);
    }

    /// A cancel that arrives after the task already finished must not take the
    /// outcome away: `get_state` reports a `result` only under `completed`, so
    /// flipping the status would lose what the caller was polling for.
    #[cfg(all(feature = "server", not(feature = "legacy-spec")))]
    #[test]
    fn a_late_cancel_leaves_a_terminal_task_alone() {
        for (finish, expected) in [(true, TaskStatus::Completed), (false, TaskStatus::Failed)] {
            let tracker = TaskTracker::new();
            let task = Task::new();
            let task_id = task.id.clone();
            let _handle = tracker.track(task);

            if finish {
                tracker.set_outcome(&task_id, Ok(serde_json::json!({ "content": [] })));
                tracker.complete(&task_id);
            } else {
                tracker.set_outcome(&task_id, Err(Error::new(ErrorCode::InternalError, "boom")));
                tracker.fail(&task_id);
            }

            let cancelled = tracker.cancel(&task_id).unwrap();
            assert_eq!(cancelled.status, expected, "cancel must not rewrite it");

            let state = tracker.get_state(&task_id).unwrap();
            assert_eq!(state.status, expected);
            assert_eq!(state.result.is_some(), finish, "the outcome survives");
            assert_eq!(state.error.is_some(), !finish);
        }
    }

    /// A cancelled task is `CancelledTask` on the wire -- status only. Neither
    /// an input request it was parked on nor an outcome its tool raced to
    /// record may leak into that shape.
    #[cfg(all(feature = "server", not(feature = "legacy-spec")))]
    #[test]
    fn get_state_hides_state_a_cancelled_task_may_not_carry() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();
        let _handle = tracker.track(task);

        let _rx = tracker
            .park_input(&task_id, "k1".into(), elicit_request())
            .expect("parked");
        tracker.require_input(&task_id);
        tracker.set_outcome(&task_id, Ok(serde_json::json!({ "content": [] })));
        tracker.cancel(&task_id).unwrap();

        let state = tracker.get_state(&task_id).unwrap();
        assert_eq!(state.status, TaskStatus::Cancelled);
        assert!(state.input_requests.is_none());
        assert!(state.result.is_none());
        assert!(state.error.is_none());
    }

    /// A failed task carries `error` -- and not an input request it never
    /// answered before failing.
    #[cfg(all(feature = "server", not(feature = "legacy-spec")))]
    #[test]
    fn get_state_hides_pending_inputs_of_a_failed_task() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();
        let _handle = tracker.track(task);

        let _rx = tracker
            .park_input(&task_id, "k1".into(), elicit_request())
            .expect("parked");
        tracker.require_input(&task_id);
        tracker.set_outcome(&task_id, Err(Error::new(ErrorCode::InternalError, "boom")));
        tracker.fail(&task_id);

        let state = tracker.get_state(&task_id).unwrap();
        assert_eq!(state.status, TaskStatus::Failed);
        assert!(state.input_requests.is_none());
        assert_eq!(state.error.expect("error")["message"], "boom");
    }

    #[cfg(all(feature = "server", not(feature = "legacy-spec")))]
    #[test]
    fn get_state_reports_the_terminal_outcome() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();
        let _handle = tracker.track(task);

        tracker.set_outcome(&task_id, Ok(serde_json::json!({ "content": [] })));
        tracker.complete(&task_id);

        let state = tracker.get_state(&task_id).unwrap();
        assert_eq!(state.status, TaskStatus::Completed);
        // The stored value is the result the request would have returned
        // synchronously, so it carries the same mandatory discriminator -- it
        // just never went through `Response::success` to get it.
        assert_eq!(
            state.result,
            Some(serde_json::json!({ "content": [], "resultType": "complete" }))
        );
        assert!(state.error.is_none());
    }

    /// A result that already says what it is keeps saying it.
    #[cfg(all(feature = "server", not(feature = "legacy-spec")))]
    #[test]
    fn set_outcome_keeps_an_existing_discriminator() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();
        let _handle = tracker.track(task);

        tracker.set_outcome(
            &task_id,
            Ok(serde_json::json!({ "resultType": "input_required" })),
        );
        tracker.complete(&task_id);

        assert_eq!(
            tracker.get_state(&task_id).unwrap().result,
            Some(serde_json::json!({ "resultType": "input_required" }))
        );
    }

    #[cfg(all(feature = "server", not(feature = "legacy-spec")))]
    #[test]
    fn get_state_reports_the_failure_error() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();
        let _handle = tracker.track(task);

        tracker.set_outcome(&task_id, Err(Error::new(ErrorCode::InternalError, "boom")));
        tracker.fail(&task_id);

        let state = tracker.get_state(&task_id).unwrap();
        assert_eq!(state.status, TaskStatus::Failed);
        assert!(state.result.is_none());
        assert_eq!(state.error.expect("error")["message"], "boom");
    }

    #[cfg(all(feature = "server", not(feature = "legacy-spec")))]
    fn elicit_request() -> crate::types::mrtr::InputRequest {
        crate::types::mrtr::InputRequest::Elicitation(
            crate::types::ElicitRequestParams::form("Sure?").into(),
        )
    }

    #[test]
    fn it_does_maintain_task_state_transitions() {
        let tracker = TaskTracker::new();
        let task = Task::new();
        let task_id = task.id.clone();

        let _handle = tracker.track(task.clone());

        let status = tracker.get_status(&task_id).unwrap();
        assert_eq!(status.status, TaskStatus::Working);

        tracker.complete(&task_id);
        let status = tracker.get_status(&task_id).unwrap();
        assert_eq!(status.status, TaskStatus::Completed);
    }

    #[test]
    fn it_does_remove_expired_completed_tasks() {
        let tracker = TaskTracker::new();
        let task = Task::from(crate::types::TaskMetadata { ttl: Some(1) });
        let task_id = task.id.clone();

        let _handle = tracker.track(task);
        tracker.complete(&task_id);

        std::thread::sleep(std::time::Duration::from_millis(10));

        assert!(tracker.get_status(&task_id).is_err());
    }
}