pe_tasks/

file.rs

//! # FileTaskRegistry - JSON-backed task storage.
//!
//! A small persistent registry for local agents. It keeps the same semantics as
//! [`InMemoryTaskRegistry`](crate::InMemoryTaskRegistry) while storing tasks and
//! dependencies in a workspace file.

use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::{Mutex, MutexGuard};

use chrono::Utc;
use pe_core::PeError;
use serde::{Deserialize, Serialize};

use crate::dependency::{self, TaskDependency};
use crate::lifecycle::{DefaultLifecycle, TaskLifecycle};
use crate::registry::{TaskFilter, TaskRegistry};
use crate::task::{Task, TaskId, TaskStatus};

#[derive(Clone, Debug, Serialize, Deserialize)]
struct TaskStoreSnapshot {
    #[serde(default = "current_version")]
    version: u32,
    #[serde(default)]
    tasks: Vec<Task>,
    #[serde(default)]
    dependencies: Vec<TaskDependency>,
}

impl Default for TaskStoreSnapshot {
    fn default() -> Self {
        Self {
            version: current_version(),
            tasks: Vec::new(),
            dependencies: Vec::new(),
        }
    }
}

fn current_version() -> u32 {
    1
}

/// JSON-file-backed [`TaskRegistry`].
///
/// This backend is meant for local app use: durable across short-lived process
/// launches, simple to inspect, and dependency-free. It is not a multi-writer
/// database; concurrent writers can still race at the file level.
pub struct FileTaskRegistry {
    path: PathBuf,
    state: Mutex<TaskStoreSnapshot>,
    lifecycle: Box<dyn TaskLifecycle>,
}

impl FileTaskRegistry {
    /// Create a registry backed by `path` with the default lifecycle.
    pub fn new(path: impl Into<PathBuf>) -> Result<Self, PeError> {
        Self::with_lifecycle(path, DefaultLifecycle)
    }

    /// Create a registry backed by `path` with a custom lifecycle.
    pub fn with_lifecycle(
        path: impl Into<PathBuf>,
        lifecycle: impl TaskLifecycle + 'static,
    ) -> Result<Self, PeError> {
        let path = path.into();
        let state = load_snapshot(&path)?;
        Ok(Self {
            path,
            state: Mutex::new(state),
            lifecycle: Box::new(lifecycle),
        })
    }

    /// Path to the JSON store.
    #[must_use]
    pub fn path(&self) -> &Path {
        &self.path
    }

    fn state_guard(&self) -> MutexGuard<'_, TaskStoreSnapshot> {
        match self.state.lock() {
            Ok(guard) => guard,
            Err(poisoned) => poisoned.into_inner(),
        }
    }

    fn persist_snapshot(&self, snapshot: &TaskStoreSnapshot) -> Result<(), PeError> {
        if let Some(parent) = self.path.parent() {
            fs::create_dir_all(parent).map_err(|error| PeError::Storage {
                details: format!("failed to create task store {}: {error}", parent.display()),
            })?;
        }
        let content = serde_json::to_string_pretty(snapshot).map_err(|error| PeError::Storage {
            details: format!(
                "failed to serialize task store {}: {error}",
                self.path.display()
            ),
        })?;
        fs::write(&self.path, content).map_err(|error| PeError::Storage {
            details: format!(
                "failed to write task store {}: {error}",
                self.path.display()
            ),
        })
    }

    fn collect_tree(root_id: &str, tasks: &[Task]) -> Vec<Task> {
        let mut result = Vec::new();
        let mut stack = vec![root_id.to_string()];
        while let Some(id) = stack.pop() {
            for task in tasks {
                if task.parent_task_id.as_deref() == Some(&id) && task.deleted_at.is_none() {
                    stack.push(task.id.clone());
                    result.push(task.clone());
                }
            }
        }
        result
    }
}

#[async_trait::async_trait]
impl TaskRegistry for FileTaskRegistry {
    async fn create(&self, task: Task) -> Result<Task, PeError> {
        {
            let mut state = self.state_guard();
            let mut next = state.clone();
            dependency::validate_parent_link(&task, &next.tasks)?;
            next.tasks.push(task.clone());
            self.persist_snapshot(&next)?;
            *state = next;
        }
        self.lifecycle.on_create(&task);
        Ok(task)
    }

    async fn get(&self, id: &TaskId) -> Result<Option<Task>, PeError> {
        let state = self.state_guard();
        Ok(state
            .tasks
            .iter()
            .find(|task| task.id == *id && task.deleted_at.is_none())
            .cloned())
    }

    async fn update(&self, task: &Task) -> Result<Task, PeError> {
        let mut updated = task.clone();
        updated.updated_at = Some(Utc::now());

        let mut state = self.state_guard();
        let mut next = state.clone();
        dependency::validate_parent_link(&updated, &next.tasks)?;
        if let Some(existing) = next
            .tasks
            .iter_mut()
            .find(|existing| existing.id == task.id)
        {
            *existing = updated.clone();
        } else {
            next.tasks.push(updated.clone());
        }
        self.persist_snapshot(&next)?;
        *state = next;
        Ok(updated)
    }

    async fn delete(&self, id: &TaskId) -> Result<bool, PeError> {
        let mut state = self.state_guard();
        let mut next = state.clone();
        let Some(task) = next.tasks.iter_mut().find(|task| task.id == *id) else {
            return Ok(false);
        };
        task.deleted_at = Some(Utc::now());
        self.persist_snapshot(&next)?;
        *state = next;
        Ok(true)
    }

    async fn restore(&self, id: &TaskId) -> Result<bool, PeError> {
        let mut state = self.state_guard();
        let mut next = state.clone();
        let Some(task) = next
            .tasks
            .iter_mut()
            .find(|task| task.id == *id && task.deleted_at.is_some())
        else {
            return Ok(false);
        };
        task.deleted_at = None;
        self.persist_snapshot(&next)?;
        *state = next;
        Ok(true)
    }

    async fn update_status(
        &self,
        id: &TaskId,
        status: TaskStatus,
        result: Option<serde_json::Value>,
        error: Option<String>,
    ) -> Result<Task, PeError> {
        let (updated, old_status) = {
            let mut state = self.state_guard();
            let mut next = state.clone();
            let task = next
                .tasks
                .iter_mut()
                .find(|task| task.id == *id)
                .ok_or_else(|| PeError::NodeNotFound { node: id.clone() })?;

            self.lifecycle.validate_transition(task, &status)?;

            let old_status = task.status.clone();
            task.status = status.clone();
            task.updated_at = Some(Utc::now());
            if let Some(result) = result {
                task.result = Some(result);
            }
            if let Some(error) = error {
                task.error = Some(error);
            }
            if status == TaskStatus::Completed {
                task.completed_at = Some(Utc::now());
            }
            let updated = task.clone();

            self.persist_snapshot(&next)?;
            *state = next;
            (updated, old_status)
        };

        self.lifecycle.on_transition(&updated, &old_status);
        match &status {
            TaskStatus::Completed => self.lifecycle.on_complete(&updated),
            TaskStatus::Failed => self
                .lifecycle
                .on_fail(&updated, updated.error.as_deref().unwrap_or("")),
            TaskStatus::Cancelled => self.lifecycle.on_cancel(&updated),
            _ => {}
        }

        Ok(updated)
    }

    async fn list(&self, filter: &TaskFilter) -> Result<Vec<Task>, PeError> {
        let state = self.state_guard();
        let limit = if filter.limit == 0 { 100 } else { filter.limit };
        let results = state
            .tasks
            .iter()
            .filter(|task| filter.include_deleted || task.deleted_at.is_none())
            .filter(|task| filter.status.as_ref().is_none_or(|s| task.status == *s))
            .filter(|task| {
                filter
                    .task_type
                    .as_ref()
                    .is_none_or(|task_type| task.task_type == *task_type)
            })
            .filter(|task| {
                filter
                    .priority
                    .as_ref()
                    .is_none_or(|priority| task.priority == *priority)
            })
            .filter(|task| {
                filter
                    .agent_id
                    .as_ref()
                    .is_none_or(|agent| task.agent_id.as_ref() == Some(agent))
            })
            .filter(|task| {
                filter
                    .assignee
                    .as_ref()
                    .is_none_or(|assignee| task.assignee == *assignee)
            })
            .filter(|task| {
                filter
                    .parent_id
                    .as_ref()
                    .is_none_or(|parent| task.parent_task_id.as_ref() == Some(parent))
            })
            .filter(|task| {
                filter
                    .tag
                    .as_ref()
                    .is_none_or(|tag| task.tags.contains(tag))
            })
            .take(limit)
            .cloned()
            .collect();
        Ok(results)
    }

    async fn get_subtasks(&self, parent_id: &TaskId) -> Result<Vec<Task>, PeError> {
        let state = self.state_guard();
        Ok(state
            .tasks
            .iter()
            .filter(|task| {
                task.parent_task_id.as_ref() == Some(parent_id) && task.deleted_at.is_none()
            })
            .cloned()
            .collect())
    }

    async fn get_tree(&self, root_id: &TaskId) -> Result<Vec<Task>, PeError> {
        let state = self.state_guard();
        Ok(Self::collect_tree(root_id, &state.tasks))
    }

    async fn add_dependency(&self, dep: TaskDependency) -> Result<(), PeError> {
        let mut state = self.state_guard();
        let mut next = state.clone();
        dependency::validate_dependency_endpoints(&dep, &next.tasks)?;
        dependency::validate_dependency_unique(&dep, &next.dependencies)?;
        if dependency::would_create_cycle(&dep.task_id, &dep.depends_on_id, &next.dependencies) {
            return Err(PeError::InvalidUpdate {
                details: format!(
                    "Adding dependency {} -> {} would create a cycle",
                    dep.task_id, dep.depends_on_id
                ),
            });
        }
        next.dependencies.push(dep);
        self.persist_snapshot(&next)?;
        *state = next;
        Ok(())
    }

    async fn remove_dependency(
        &self,
        task_id: &TaskId,
        depends_on_id: &TaskId,
    ) -> Result<(), PeError> {
        let mut state = self.state_guard();
        let mut next = state.clone();
        next.dependencies
            .retain(|dep| !(dep.task_id == *task_id && dep.depends_on_id == *depends_on_id));
        self.persist_snapshot(&next)?;
        *state = next;
        Ok(())
    }

    async fn get_dependencies(&self, task_id: &TaskId) -> Result<Vec<TaskDependency>, PeError> {
        let state = self.state_guard();
        Ok(state
            .dependencies
            .iter()
            .filter(|dep| dep.task_id == *task_id)
            .cloned()
            .collect())
    }

    async fn get_dependents(&self, task_id: &TaskId) -> Result<Vec<TaskDependency>, PeError> {
        let state = self.state_guard();
        Ok(state
            .dependencies
            .iter()
            .filter(|dep| dep.depends_on_id == *task_id)
            .cloned()
            .collect())
    }

    async fn get_ready_tasks(&self) -> Result<Vec<Task>, PeError> {
        let state = self.state_guard();
        let pending = state
            .tasks
            .iter()
            .filter(|task| task.status == TaskStatus::Pending && task.deleted_at.is_none())
            .map(|task| task.id.clone())
            .collect::<Vec<_>>();
        let statuses: HashMap<TaskId, TaskStatus> = state
            .tasks
            .iter()
            .map(|task| (task.id.clone(), task.status.clone()))
            .collect();
        let ready_ids = dependency::find_ready_tasks(&pending, &state.dependencies, &statuses);
        Ok(ready_ids
            .iter()
            .filter_map(|id| state.tasks.iter().find(|task| task.id == *id).cloned())
            .collect())
    }
}

fn load_snapshot(path: &Path) -> Result<TaskStoreSnapshot, PeError> {
    let content = match fs::read_to_string(path) {
        Ok(content) => content,
        Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
            return Ok(TaskStoreSnapshot::default());
        }
        Err(error) => {
            return Err(PeError::Storage {
                details: format!("failed to read task store {}: {error}", path.display()),
            });
        }
    };
    if content.trim().is_empty() {
        return Ok(TaskStoreSnapshot::default());
    }
    serde_json::from_str(&content).map_err(|error| PeError::Storage {
        details: format!("failed to parse task store {}: {error}", path.display()),
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::dependency::DependencyType;

    fn temp_path(name: &str) -> PathBuf {
        std::env::temp_dir().join(format!("pe-tasks-{name}-{}.json", uuid::Uuid::new_v4()))
    }

    #[tokio::test]
    async fn persists_tasks_between_instances() {
        let path = temp_path("persist");
        let registry = FileTaskRegistry::new(&path).unwrap();
        let created = registry
            .create(Task::agent_task("Persistent task", "agent-1"))
            .await
            .unwrap();
        drop(registry);

        let reloaded = FileTaskRegistry::new(&path).unwrap();
        let task = reloaded.get(&created.id).await.unwrap().unwrap();
        assert_eq!(task.title, "Persistent task");
        assert_eq!(task.agent_id.as_deref(), Some("agent-1"));

        let _ = fs::remove_file(path);
    }

    #[tokio::test]
    async fn persists_dependencies_between_instances() {
        let path = temp_path("deps");
        let registry = FileTaskRegistry::new(&path).unwrap();
        let blocker = registry.create(Task::new("Blocker")).await.unwrap();
        let blocked = registry.create(Task::new("Blocked")).await.unwrap();
        registry
            .add_dependency(TaskDependency::new(
                &blocked.id,
                &blocker.id,
                DependencyType::Blocks,
            ))
            .await
            .unwrap();
        drop(registry);

        let reloaded = FileTaskRegistry::new(&path).unwrap();
        let deps = reloaded.get_dependencies(&blocked.id).await.unwrap();
        assert_eq!(deps.len(), 1);
        assert_eq!(deps[0].depends_on_id, blocker.id);

        let _ = fs::remove_file(path);
    }

    #[tokio::test]
    async fn validates_dependency_endpoints_and_duplicates() {
        let path = temp_path("dep-validation");
        let registry = FileTaskRegistry::new(&path).unwrap();
        let blocker = registry.create(Task::new("Blocker")).await.unwrap();
        let blocked = registry.create(Task::new("Blocked")).await.unwrap();

        let err = registry
            .add_dependency(TaskDependency::new(
                &blocked.id,
                "missing",
                DependencyType::Blocks,
            ))
            .await
            .unwrap_err();
        assert!(matches!(err, PeError::NodeNotFound { node } if node == "missing"));

        registry
            .add_dependency(TaskDependency::new(
                &blocked.id,
                &blocker.id,
                DependencyType::Blocks,
            ))
            .await
            .unwrap();

        let err = registry
            .add_dependency(TaskDependency::new(
                &blocked.id,
                &blocker.id,
                DependencyType::Related,
            ))
            .await
            .unwrap_err();
        assert!(
            matches!(err, PeError::InvalidUpdate { details } if details.contains("already exists"))
        );

        let _ = fs::remove_file(path);
    }

    #[tokio::test]
    async fn validates_parent_links() {
        let path = temp_path("parent-validation");
        let registry = FileTaskRegistry::new(&path).unwrap();

        let err = registry
            .create(Task::new("Orphan").with_parent("missing-parent"))
            .await
            .unwrap_err();
        assert!(matches!(err, PeError::NodeNotFound { node } if node == "missing-parent"));

        let parent = registry.create(Task::new("Parent")).await.unwrap();
        let child = registry
            .create(Task::new("Child").with_parent(&parent.id))
            .await
            .unwrap();
        let mut updated_parent = parent.clone();
        updated_parent.parent_task_id = Some(child.id.clone());

        let err = registry.update(&updated_parent).await.unwrap_err();
        assert!(matches!(err, PeError::InvalidUpdate { .. }));

        let stored_parent = registry.get(&parent.id).await.unwrap().unwrap();
        assert!(stored_parent.parent_task_id.is_none());

        let _ = fs::remove_file(path);
    }
}