mcp_host/managers/

task.rs

//! Task management for async task execution
//!
//! Provides storage and lifecycle management for MCP tasks.

use std::collections::HashSet;
use std::sync::Arc;
use std::time::{Duration, Instant};

use chrono::Utc;
use dashmap::DashMap;
use serde_json::Value;
use thiserror::Error;
use tokio::sync::{RwLock, oneshot};
use uuid::Uuid;

use crate::protocol::types::{JsonRpcRequest, Task, TaskStatus};

/// Task-related errors
#[derive(Debug, Error)]
pub enum TaskError {
    /// Task not found
    #[error("Task not found: {0}")]
    NotFound(String),

    /// Task timeout
    #[error("Task timeout after {0:?}")]
    Timeout(Duration),

    /// Task failed
    #[error("Task failed: {0}")]
    Failed(String),

    /// Task already in terminal state
    #[error("Task already in terminal status: {0:?}")]
    AlreadyTerminal(TaskStatus),

    /// Internal error
    #[error("Internal error: {0}")]
    Internal(String),
}

/// Internal task entry with result channel
struct TaskEntry {
    /// Public task metadata
    task: Arc<RwLock<Task>>,
    /// Session ID owning this task
    session_id: String,
    /// Original request that created the task (reserved for future use)
    #[allow(dead_code)]
    original_request: JsonRpcRequest,
    /// Channel for waiting on result
    result_tx: Option<oneshot::Sender<Value>>,
    /// Stored result (when completed)
    result: Arc<RwLock<Option<Value>>>,
    /// Expiration time
    expires_at: Option<Instant>,
}

/// Task store for managing async task execution
pub struct TaskStore {
    /// All tasks indexed by task ID
    tasks: DashMap<String, TaskEntry>,
    /// Session → task IDs index
    by_session: DashMap<String, HashSet<String>>,
    /// Default TTL for new tasks
    default_ttl: Duration,
    /// Default poll interval
    default_poll_interval: Duration,
}

impl TaskStore {
    /// Create a new task store
    pub fn new(default_ttl: Duration, default_poll_interval: Duration) -> Self {
        Self {
            tasks: DashMap::new(),
            by_session: DashMap::new(),
            default_ttl,
            default_poll_interval,
        }
    }

    /// Create a new task
    pub fn create_task(
        &self,
        session_id: &str,
        original_request: JsonRpcRequest,
        requested_ttl: Option<u64>,
    ) -> (Task, oneshot::Receiver<Value>) {
        let task_id = Uuid::new_v4().to_string();
        let now = Utc::now().to_rfc3339();

        // Determine TTL (use requested or default)
        let ttl_duration = requested_ttl
            .map(Duration::from_millis)
            .unwrap_or(self.default_ttl);
        let ttl_ms = if ttl_duration == Duration::ZERO {
            None
        } else {
            Some(ttl_duration.as_millis() as u64)
        };

        let task = Task {
            task_id: task_id.clone(),
            status: TaskStatus::Working,
            status_message: None,
            created_at: now.clone(),
            last_updated_at: now,
            ttl: ttl_ms,
            poll_interval: Some(self.default_poll_interval.as_millis() as u64),
        };

        // Create result channel
        let (result_tx, result_rx) = oneshot::channel();

        let entry = TaskEntry {
            task: Arc::new(RwLock::new(task.clone())),
            session_id: session_id.to_string(),
            original_request,
            result_tx: Some(result_tx),
            result: Arc::new(RwLock::new(None)),
            expires_at: ttl_ms.map(|ms| Instant::now() + Duration::from_millis(ms)),
        };

        // Store task
        self.tasks.insert(task_id.clone(), entry);

        // Index by session
        self.by_session
            .entry(session_id.to_string())
            .or_default()
            .insert(task_id);

        (task, result_rx)
    }

    /// Get task by ID
    pub async fn get_task(&self, task_id: &str) -> Option<Task> {
        let entry = self.tasks.get(task_id)?;
        Some(entry.task.read().await.clone())
    }

    /// Get task for specific session (enforces session isolation)
    pub async fn get_task_for_session(&self, task_id: &str, session_id: &str) -> Option<Task> {
        let entry = self.tasks.get(task_id)?;
        if entry.session_id == session_id {
            Some(entry.task.read().await.clone())
        } else {
            None
        }
    }

    /// List tasks for a session (with pagination)
    pub async fn list_tasks(
        &self,
        session_id: &str,
        cursor: Option<&str>,
        limit: usize,
    ) -> (Vec<Task>, Option<String>) {
        let task_ids = match self.by_session.get(session_id) {
            Some(ids) => ids.clone(),
            None => return (vec![], None),
        };

        let mut task_ids: Vec<_> = task_ids.into_iter().collect();
        task_ids.sort(); // Deterministic ordering

        // Find cursor position
        let start = if let Some(c) = cursor {
            task_ids
                .iter()
                .position(|id| id == c)
                .map(|p| p + 1)
                .unwrap_or(0)
        } else {
            0
        };

        let end = (start + limit).min(task_ids.len());
        let page_ids = &task_ids[start..end];

        // Collect tasks
        let mut tasks = Vec::new();
        for id in page_ids {
            if let Some(task) = self.get_task(id).await {
                tasks.push(task);
            }
        }

        // Next cursor
        let next_cursor = if end < task_ids.len() {
            task_ids.get(end).cloned()
        } else {
            None
        };

        (tasks, next_cursor)
    }

    /// Update task status
    pub async fn update_status(
        &self,
        task_id: &str,
        new_status: TaskStatus,
        status_message: Option<String>,
    ) -> Result<(), TaskError> {
        let entry = self
            .tasks
            .get(task_id)
            .ok_or_else(|| TaskError::NotFound(task_id.to_string()))?;

        let mut task = entry.task.write().await;

        // Check valid transition
        if let (TaskStatus::Completed | TaskStatus::Failed | TaskStatus::Cancelled, _) =
            (&task.status, &new_status)
        {
            return Err(TaskError::AlreadyTerminal(task.status));
        }

        task.status = new_status;
        task.status_message = status_message;
        task.last_updated_at = Utc::now().to_rfc3339();

        Ok(())
    }

    /// Store task result
    pub async fn store_result(&self, task_id: &str, result: Value) -> Result<(), TaskError> {
        let entry = self
            .tasks
            .get(task_id)
            .ok_or_else(|| TaskError::NotFound(task_id.to_string()))?;

        // Store result
        *entry.result.write().await = Some(result.clone());

        // Send to waiting receiver (if any)
        drop(entry); // Drop the read lock before taking result_tx
        if let Some(mut entry_mut) = self.tasks.get_mut(task_id)
            && let Some(tx) = entry_mut.result_tx.take()
        {
            let _ = tx.send(result);
        }

        Ok(())
    }

    /// Get stored result
    pub async fn get_result(&self, task_id: &str) -> Option<Value> {
        let entry = self.tasks.get(task_id)?;
        entry.result.read().await.clone()
    }

    /// Wait for task result (blocks until terminal state)
    pub async fn wait_for_result(
        &self,
        task_id: &str,
        timeout: Duration,
    ) -> Result<Value, TaskError> {
        let entry = self
            .tasks
            .get(task_id)
            .ok_or_else(|| TaskError::NotFound(task_id.to_string()))?;

        let task = entry.task.read().await.clone();

        // If already terminal, return stored result
        if matches!(
            task.status,
            TaskStatus::Completed | TaskStatus::Failed | TaskStatus::Cancelled
        ) {
            drop(task);
            if let Some(result) = self.get_result(task_id).await {
                return Ok(result);
            } else {
                return Err(TaskError::Failed(
                    "Task in terminal state but no result stored".to_string(),
                ));
            }
        }

        // Otherwise wait on channel
        drop(entry);

        // Get a new receiver (or clone existing one)
        // Note: this is a simplification - in production you'd want a broadcast channel
        // For now, we'll poll
        let start = Instant::now();
        loop {
            if start.elapsed() > timeout {
                return Err(TaskError::Timeout(timeout));
            }

            let task = self
                .get_task(task_id)
                .await
                .ok_or_else(|| TaskError::NotFound(task_id.to_string()))?;

            if matches!(
                task.status,
                TaskStatus::Completed | TaskStatus::Failed | TaskStatus::Cancelled
            ) {
                if let Some(result) = self.get_result(task_id).await {
                    return Ok(result);
                } else {
                    return Err(TaskError::Failed(
                        "Task completed but no result stored".to_string(),
                    ));
                }
            }

            tokio::time::sleep(Duration::from_millis(100)).await;
        }
    }

    /// Cancel a task
    pub async fn cancel_task(&self, task_id: &str, session_id: &str) -> Result<Task, TaskError> {
        // Verify session owns task
        let entry = self
            .tasks
            .get(task_id)
            .ok_or_else(|| TaskError::NotFound(task_id.to_string()))?;

        if entry.session_id != session_id {
            return Err(TaskError::NotFound(task_id.to_string()));
        }

        let mut task = entry.task.write().await;

        // Check not already terminal
        if matches!(
            task.status,
            TaskStatus::Completed | TaskStatus::Failed | TaskStatus::Cancelled
        ) {
            return Err(TaskError::AlreadyTerminal(task.status));
        }

        // Update to cancelled
        task.status = TaskStatus::Cancelled;
        task.status_message = Some("Cancelled by request".to_string());
        task.last_updated_at = Utc::now().to_rfc3339();

        Ok(task.clone())
    }

    /// Cleanup expired tasks (should be called periodically)
    pub async fn cleanup_expired(&self) {
        let now = Instant::now();
        let mut to_remove = Vec::new();

        for entry in self.tasks.iter() {
            if let Some(expires_at) = entry.expires_at
                && now >= expires_at
            {
                to_remove.push(entry.key().clone());
            }
        }

        for task_id in to_remove {
            if let Some((_, entry)) = self.tasks.remove(&task_id) {
                // Remove from session index
                if let Some(mut ids) = self.by_session.get_mut(&entry.session_id) {
                    ids.remove(&task_id);
                }
            }
        }
    }

    /// Spawn background cleanup task
    pub fn spawn_cleanup_task(self: Arc<Self>, interval: Duration) {
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(interval);
            loop {
                interval.tick().await;
                self.cleanup_expired().await;
            }
        });
    }
}

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

    #[tokio::test]
    async fn test_create_task() {
        let store = TaskStore::new(Duration::from_secs(60), Duration::from_secs(5));
        let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));

        let (task, _rx) = store.create_task("session1", req, None);

        assert_eq!(task.status, TaskStatus::Working);
        assert!(task.task_id.len() > 0);
        assert_eq!(task.ttl, Some(60000));
    }

    #[tokio::test]
    async fn test_get_task() {
        let store = TaskStore::new(Duration::from_secs(60), Duration::from_secs(5));
        let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));

        let (task, _rx) = store.create_task("session1", req, None);
        let retrieved = store.get_task(&task.task_id).await.unwrap();

        assert_eq!(retrieved.task_id, task.task_id);
        assert_eq!(retrieved.status, TaskStatus::Working);
    }

    #[tokio::test]
    async fn test_session_isolation() {
        let store = TaskStore::new(Duration::from_secs(60), Duration::from_secs(5));
        let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));

        let (task, _rx) = store.create_task("session1", req, None);

        // Should not be visible from different session
        assert!(
            store
                .get_task_for_session(&task.task_id, "session2")
                .await
                .is_none()
        );

        // Should be visible from same session
        assert!(
            store
                .get_task_for_session(&task.task_id, "session1")
                .await
                .is_some()
        );
    }

    #[tokio::test]
    async fn test_update_status() {
        let store = TaskStore::new(Duration::from_secs(60), Duration::from_secs(5));
        let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));

        let (task, _rx) = store.create_task("session1", req, None);

        store
            .update_status(
                &task.task_id,
                TaskStatus::Completed,
                Some("Done".to_string()),
            )
            .await
            .unwrap();

        let updated = store.get_task(&task.task_id).await.unwrap();
        assert_eq!(updated.status, TaskStatus::Completed);
        assert_eq!(updated.status_message, Some("Done".to_string()));
    }

    #[tokio::test]
    async fn test_store_and_get_result() {
        let store = TaskStore::new(Duration::from_secs(60), Duration::from_secs(5));
        let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));

        let (task, _rx) = store.create_task("session1", req, None);

        let result = json!({"answer": 42});
        store
            .store_result(&task.task_id, result.clone())
            .await
            .unwrap();

        let retrieved = store.get_result(&task.task_id).await.unwrap();
        assert_eq!(retrieved, result);
    }

    #[tokio::test]
    async fn test_cancel_task() {
        let store = TaskStore::new(Duration::from_secs(60), Duration::from_secs(5));
        let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));

        let (task, _rx) = store.create_task("session1", req, None);

        let cancelled = store.cancel_task(&task.task_id, "session1").await.unwrap();
        assert_eq!(cancelled.status, TaskStatus::Cancelled);
    }

    #[tokio::test]
    async fn test_list_tasks() {
        let store = TaskStore::new(Duration::from_secs(60), Duration::from_secs(5));

        // Create multiple tasks
        for _ in 0..5 {
            let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));
            store.create_task("session1", req, None);
        }

        let (tasks, cursor) = store.list_tasks("session1", None, 10).await;
        assert_eq!(tasks.len(), 5);
        assert!(cursor.is_none());
    }
}