ai-dispatch 8.99.6

// Extra workgroup persistence for lifecycle commands.
// Extends Store with update/delete flows while keeping store.rs smaller.
// Depends on rusqlite, chrono, and the core Store type.

use anyhow::Result;
use chrono::Local;
use rusqlite::params;

use crate::store::Store;

impl Store {
    pub fn update_workgroup(
        &self,
        id: &str,
        name: Option<&str>,
        shared_context: Option<&str>,
    ) -> Result<Option<crate::types::Workgroup>> {
        let Some(mut workgroup) = self.get_workgroup(id)? else {
            return Ok(None);
        };
        if let Some(name) = name {
            workgroup.name = name.to_string();
        }
        if let Some(shared_context) = shared_context {
            workgroup.shared_context = shared_context.to_string();
        }
        workgroup.updated_at = Local::now();

        self.db().execute(
            "UPDATE workgroups SET name = ?1, shared_context = ?2, updated_at = ?3
             WHERE id = ?4",
            params![
                workgroup.name,
                workgroup.shared_context,
                workgroup.updated_at.to_rfc3339(),
                workgroup.id.as_str(),
            ],
        )?;

        Ok(Some(workgroup))
    }

    pub fn delete_workgroup(&self, id: &str) -> Result<Option<usize>> {
        let mut conn = self.db();
        let tx = conn.transaction()?;
        let tagged_tasks = Self::count_workgroup_tasks_in_tx(&tx, id)?;
        let deleted = tx.execute("DELETE FROM workgroups WHERE id = ?1", params![id])?;
        if deleted == 0 {
            return Ok(None);
        }
        tx.commit()?;
        drop(conn);
        remove_workspace_dir(id)?;
        Ok(Some(tagged_tasks))
    }

    pub fn delete_workgroup_cascade(&self, id: &str) -> Result<Option<usize>> {
        let mut conn = self.db();
        let tx = conn.transaction()?;
        tx.execute(
            "DELETE FROM events
             WHERE task_id IN (SELECT id FROM tasks WHERE workgroup_id = ?1)",
            params![id],
        )?;
        let deleted_tasks = tx.execute("DELETE FROM tasks WHERE workgroup_id = ?1", params![id])?;
        let deleted = tx.execute("DELETE FROM workgroups WHERE id = ?1", params![id])?;
        if deleted == 0 {
            return Ok(None);
        }
        tx.commit()?;
        drop(conn);
        remove_workspace_dir(id)?;
        Ok(Some(deleted_tasks))
    }

    fn count_workgroup_tasks_in_tx(tx: &rusqlite::Transaction<'_>, id: &str) -> Result<usize> {
        let count = tx.query_row(
            "SELECT COUNT(*) FROM tasks WHERE workgroup_id = ?1",
            params![id],
            |row| row.get::<_, i64>(0),
        )?;
        Ok(count as usize)
    }
}

fn remove_workspace_dir(id: &str) -> Result<()> {
    let workspace_dir = crate::paths::workspace_dir(id)?;
    let ws = workspace_dir.to_string_lossy();
    let safe_prefix =
        ws.starts_with("/tmp/aid-wg-") || ws.starts_with("/private/tmp/aid-wg-");
    // In tests, AidHomeGuard reroutes workspace_dir under a TempDir — that path
    // is isolated per-test, so it's safe to remove even though it lacks the
    // /tmp/aid-wg-* prefix.
    let test_override_active = cfg!(test)
        && ws.contains("/workgroups/");
    if safe_prefix || test_override_active {
        let _ = std::fs::remove_dir_all(&workspace_dir);
    } else {
        aid_warn!(
            "[aid] SAFETY: refusing to remove workspace '{}' — not under /tmp/aid-wg-*",
            ws
        );
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{AgentKind, Task, TaskId, TaskStatus, VerifyStatus};
    use chrono::Local;

    fn make_task(id: &str, group_id: &str) -> Task {
        Task {
            id: TaskId(id.to_string()),
            agent: AgentKind::Codex,
            custom_agent_name: None,
            prompt: "test prompt".to_string(),
            resolved_prompt: None,
            category: None,
            status: TaskStatus::Done,
            parent_task_id: None,
            workgroup_id: Some(group_id.to_string()),
            caller_kind: None,
            caller_session_id: None,
            agent_session_id: None,
            repo_path: None,
            worktree_path: None,
            worktree_branch: None,
            start_sha: None,
            log_path: None,
            output_path: None,
            tokens: None,
            prompt_tokens: None,
            duration_ms: None,
            model: None,
            cost_usd: None,
            exit_code: None,
            created_at: Local::now(),
            completed_at: None,
            verify: None,
            verify_status: VerifyStatus::Skipped,
            pending_reason: None,
            read_only: false,
            budget: false,
            audit_verdict: None,
            audit_report_path: None,
            delivery_assessment: None,
        }
    }

    #[test]
    fn update_workgroup_changes_requested_fields() {
        let store = Store::open_memory().unwrap();
        let workgroup = store
            .create_workgroup("dispatch", "Shared repo rules.", None, None)
            .unwrap();

        let updated = store
            .update_workgroup(
                workgroup.id.as_str(),
                Some("dispatch-core"),
                Some("Updated shared rules."),
            )
            .unwrap()
            .unwrap();

        assert_eq!(updated.name, "dispatch-core");
        assert_eq!(updated.shared_context, "Updated shared rules.");
        assert!(updated.updated_at >= updated.created_at);
    }

    #[test]
    fn delete_workgroup_keeps_historical_task_tags() {
        let store = Store::open_memory().unwrap();
        let workgroup = store
            .create_workgroup("dispatch", "Shared repo rules.", None, None)
            .unwrap();
        store
            .insert_task(&make_task("t-1000", workgroup.id.as_str()))
            .unwrap();

        let tagged_tasks = store.delete_workgroup(workgroup.id.as_str()).unwrap();
        let task = store.get_task("t-1000").unwrap().unwrap();

        assert_eq!(tagged_tasks, Some(1));
        assert_eq!(task.workgroup_id.as_deref(), Some(workgroup.id.as_str()));
        assert!(store
            .get_workgroup(workgroup.id.as_str())
            .unwrap()
            .is_none());
    }

    #[test]
    fn delete_workgroup_cascade_removes_group_tasks_and_events() {
        let store = Store::open_memory().unwrap();
        let workgroup = store
            .create_workgroup("dispatch", "Shared repo rules.", None, None)
            .unwrap();
        store
            .insert_task(&make_task("t-1001", workgroup.id.as_str()))
            .unwrap();
        store
            .insert_event(&crate::types::TaskEvent {
                task_id: crate::types::TaskId("t-1001".to_string()),
                timestamp: Local::now(),
                event_kind: crate::types::EventKind::Milestone,
                detail: "deleted with task".to_string(),
                metadata: None,
            })
            .unwrap();

        let deleted_tasks = store.delete_workgroup_cascade(workgroup.id.as_str()).unwrap();

        assert_eq!(deleted_tasks, Some(1));
        assert!(store.get_task("t-1001").unwrap().is_none());
        assert!(store
            .get_workgroup(workgroup.id.as_str())
            .unwrap()
            .is_none());
    }
}