tmai_core/api/

queries.rs

//! Read-only query methods on [`TmaiCore`].
//!
//! Every method acquires a read lock internally, converts to owned snapshots,
//! and releases the lock before returning. Callers never hold a lock.

use super::core::TmaiCore;
use super::types::{AgentDefinitionInfo, AgentSnapshot, ApiError, TeamSummary, TeamTaskInfo};

impl TmaiCore {
    // =========================================================
    // Agent queries
    // =========================================================

    /// List all monitored agents as owned snapshots, in current display order.
    pub fn list_agents(&self) -> Vec<AgentSnapshot> {
        let state = self.state().read();
        let defs = &state.agent_definitions;
        state
            .agent_order
            .iter()
            .filter_map(|id| state.agents.get(id))
            .map(|a| {
                let mut snap = AgentSnapshot::from_agent(a);
                snap.agent_definition = Self::match_agent_definition(a, defs);
                snap
            })
            .collect()
    }

    /// Get a single agent snapshot by target ID.
    pub fn get_agent(&self, target: &str) -> Result<AgentSnapshot, ApiError> {
        let state = self.state().read();
        let defs = &state.agent_definitions;
        state
            .agents
            .get(target)
            .map(|a| {
                let mut snap = AgentSnapshot::from_agent(a);
                snap.agent_definition = Self::match_agent_definition(a, defs);
                snap
            })
            .ok_or_else(|| ApiError::AgentNotFound {
                target: target.to_string(),
            })
    }

    /// Get the currently selected agent snapshot.
    pub fn selected_agent(&self) -> Result<AgentSnapshot, ApiError> {
        let state = self.state().read();
        let defs = &state.agent_definitions;
        state
            .selected_agent()
            .map(|agent| {
                let mut snapshot = AgentSnapshot::from_agent(agent);
                snapshot.agent_definition = Self::match_agent_definition(agent, defs);
                snapshot
            })
            .ok_or(ApiError::NoSelection)
    }

    /// Get the number of agents that need user attention.
    pub fn attention_count(&self) -> usize {
        let state = self.state().read();
        state.attention_count()
    }

    /// Get the total number of monitored agents.
    pub fn agent_count(&self) -> usize {
        let state = self.state().read();
        state.agents.len()
    }

    /// List agents that need attention (awaiting approval or error).
    pub fn agents_needing_attention(&self) -> Vec<AgentSnapshot> {
        let state = self.state().read();
        state
            .agent_order
            .iter()
            .filter_map(|id| state.agents.get(id))
            .filter(|a| a.status.needs_attention())
            .map(AgentSnapshot::from_agent)
            .collect()
    }

    // =========================================================
    // Preview
    // =========================================================

    /// Get the ANSI preview content for an agent.
    pub fn get_preview(&self, target: &str) -> Result<String, ApiError> {
        let state = self.state().read();
        state
            .agents
            .get(target)
            .map(|a| a.last_content_ansi.clone())
            .ok_or_else(|| ApiError::AgentNotFound {
                target: target.to_string(),
            })
    }

    /// Get the plain-text content for an agent.
    pub fn get_content(&self, target: &str) -> Result<String, ApiError> {
        let state = self.state().read();
        state
            .agents
            .get(target)
            .map(|a| a.last_content.clone())
            .ok_or_else(|| ApiError::AgentNotFound {
                target: target.to_string(),
            })
    }

    // =========================================================
    // Transcript queries
    // =========================================================

    /// Get transcript records for an agent (used for hybrid scrollback preview).
    ///
    /// Returns parsed JSONL records from the agent's Claude Code conversation log.
    /// The records are looked up by pane_id from the transcript registry.
    pub fn get_transcript(
        &self,
        target: &str,
    ) -> Result<Vec<crate::transcript::TranscriptRecord>, ApiError> {
        // Verify agent exists and get pane_id
        let pane_id = {
            let state = self.state().read();
            let agent = state
                .agents
                .get(target)
                .ok_or_else(|| ApiError::AgentNotFound {
                    target: target.to_string(),
                })?;
            // Use target_to_pane_id mapping, or fall back to using the target itself
            state
                .target_to_pane_id
                .get(&agent.id)
                .cloned()
                .unwrap_or_else(|| agent.id.clone())
        };

        // Look up transcript records from the registry
        let registry = match self.transcript_registry() {
            Some(reg) => reg,
            None => return Ok(Vec::new()),
        };

        let reg = registry.read();
        Ok(reg
            .get(&pane_id)
            .map(|state| state.recent_records.clone())
            .unwrap_or_default())
    }

    // =========================================================
    // Team queries
    // =========================================================

    /// List all known teams as owned summaries.
    pub fn list_teams(&self) -> Vec<TeamSummary> {
        let state = self.state().read();
        let mut teams: Vec<TeamSummary> = state
            .teams
            .values()
            .map(TeamSummary::from_snapshot)
            .collect();
        teams.sort_by(|a, b| a.name.cmp(&b.name));
        teams
    }

    /// Get a single team summary by name.
    pub fn get_team(&self, name: &str) -> Result<TeamSummary, ApiError> {
        let state = self.state().read();
        state
            .teams
            .get(name)
            .map(TeamSummary::from_snapshot)
            .ok_or_else(|| ApiError::TeamNotFound {
                name: name.to_string(),
            })
    }

    /// Get tasks for a team.
    pub fn get_team_tasks(&self, name: &str) -> Result<Vec<TeamTaskInfo>, ApiError> {
        let state = self.state().read();
        state
            .teams
            .get(name)
            .map(|ts| ts.tasks.iter().map(TeamTaskInfo::from_task).collect())
            .ok_or_else(|| ApiError::TeamNotFound {
                name: name.to_string(),
            })
    }

    // =========================================================
    // Security queries
    // =========================================================

    /// Run a config audit and cache the result in state.
    ///
    /// Acquires a read lock to gather project directories, releases it,
    /// runs the audit (no lock held), then acquires a write lock to store the result.
    pub fn config_audit(&self) -> crate::security::ScanResult {
        // Gather project directories from agent working_dir fields
        let dirs: Vec<std::path::PathBuf> = {
            let state = self.state().read();
            state
                .agents
                .values()
                .map(|a| std::path::PathBuf::from(&a.cwd))
                .collect()
        };

        // Run audit without holding any lock
        let result = crate::security::ConfigAuditScanner::scan(&dirs);

        // Store result
        {
            let mut state = self.state().write();
            state.config_audit = Some(result.clone());
        }

        result
    }

    /// Get the last cached config audit result (no new audit).
    pub fn last_config_audit(&self) -> Option<crate::security::ScanResult> {
        let state = self.state().read();
        state.config_audit.clone()
    }

    // =========================================================
    // Miscellaneous queries
    // =========================================================

    /// Match an agent to its definition by configured agent_type or member name.
    fn match_agent_definition(
        agent: &crate::agents::MonitoredAgent,
        defs: &[crate::teams::AgentDefinition],
    ) -> Option<AgentDefinitionInfo> {
        if defs.is_empty() {
            return None;
        }
        if let Some(ref team_info) = agent.team_info {
            // 1) Try configured agent_type (explicit mapping from team config)
            if let Some(ref agent_type) = team_info.agent_type {
                if let Some(def) = defs.iter().find(|d| d.name == *agent_type) {
                    return Some(AgentDefinitionInfo::from_definition(def));
                }
            }
            // 2) Fallback: try member_name as agent definition name
            if let Some(def) = defs.iter().find(|d| d.name == team_info.member_name) {
                return Some(AgentDefinitionInfo::from_definition(def));
            }
        }
        None
    }

    /// Check if the application is still running.
    pub fn is_running(&self) -> bool {
        let state = self.state().read();
        state.running
    }

    /// Get the last poll timestamp.
    pub fn last_poll(&self) -> Option<chrono::DateTime<chrono::Utc>> {
        let state = self.state().read();
        state.last_poll
    }

    /// Get known working directories from current agents.
    pub fn known_directories(&self) -> Vec<String> {
        let state = self.state().read();
        state.get_known_directories()
    }

    // =========================================================
    // Project queries
    // =========================================================

    /// List registered project directories.
    pub fn list_projects(&self) -> Vec<String> {
        let state = self.state().read();
        state.registered_projects.clone()
    }

    /// Add a project directory. Persists to config.toml.
    pub fn add_project(&self, path: &str) -> Result<(), ApiError> {
        let canonical = std::path::Path::new(path);
        if !canonical.is_absolute() {
            return Err(ApiError::InvalidInput {
                message: "Project path must be absolute".to_string(),
            });
        }
        if !canonical.is_dir() {
            return Err(ApiError::InvalidInput {
                message: format!("Directory does not exist: {}", path),
            });
        }
        let canonical_str = canonical.to_string_lossy().to_string();

        let mut state = self.state().write();
        if state.registered_projects.contains(&canonical_str) {
            return Ok(()); // Already registered, idempotent
        }
        state.registered_projects.push(canonical_str);
        let projects = state.registered_projects.clone();
        drop(state);

        crate::config::Settings::save_projects(&projects);
        Ok(())
    }

    /// Remove a project directory. Persists to config.toml.
    pub fn remove_project(&self, path: &str) -> Result<(), ApiError> {
        let mut state = self.state().write();
        let before = state.registered_projects.len();
        state.registered_projects.retain(|p| p != path);
        if state.registered_projects.len() == before {
            return Err(ApiError::InvalidInput {
                message: format!("Project not found: {}", path),
            });
        }
        let projects = state.registered_projects.clone();
        drop(state);

        crate::config::Settings::save_projects(&projects);
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::agents::{AgentStatus, AgentType, MonitoredAgent};
    use crate::api::builder::TmaiCoreBuilder;
    use crate::config::Settings;
    use crate::state::AppState;

    fn make_core_with_agents(agents: Vec<MonitoredAgent>) -> TmaiCore {
        let state = AppState::shared();
        {
            let mut s = state.write();
            s.update_agents(agents);
        }
        TmaiCoreBuilder::new(Settings::default())
            .with_state(state)
            .build()
    }

    fn test_agent(id: &str, status: AgentStatus) -> MonitoredAgent {
        let mut agent = MonitoredAgent::new(
            id.to_string(),
            AgentType::ClaudeCode,
            "Title".to_string(),
            "/home/user".to_string(),
            100,
            "main".to_string(),
            "win".to_string(),
            0,
            0,
        );
        agent.status = status;
        agent
    }

    #[test]
    fn test_list_agents_empty() {
        let core = TmaiCoreBuilder::new(Settings::default()).build();
        assert!(core.list_agents().is_empty());
    }

    #[test]
    fn test_list_agents() {
        let core = make_core_with_agents(vec![
            test_agent("main:0.0", AgentStatus::Idle),
            test_agent(
                "main:0.1",
                AgentStatus::Processing {
                    activity: "Bash".to_string(),
                },
            ),
        ]);

        let agents = core.list_agents();
        assert_eq!(agents.len(), 2);
    }

    #[test]
    fn test_get_agent_found() {
        let core = make_core_with_agents(vec![test_agent("main:0.0", AgentStatus::Idle)]);

        let result = core.get_agent("main:0.0");
        assert!(result.is_ok());
        assert_eq!(result.unwrap().id, "main:0.0");
    }

    #[test]
    fn test_get_agent_not_found() {
        let core = TmaiCoreBuilder::new(Settings::default()).build();
        let result = core.get_agent("nonexistent");
        assert!(matches!(result, Err(ApiError::AgentNotFound { .. })));
    }

    #[test]
    fn test_attention_count() {
        let core = make_core_with_agents(vec![
            test_agent("main:0.0", AgentStatus::Idle),
            test_agent(
                "main:0.1",
                AgentStatus::AwaitingApproval {
                    approval_type: crate::agents::ApprovalType::ShellCommand,
                    details: "rm -rf".to_string(),
                },
            ),
            test_agent(
                "main:0.2",
                AgentStatus::Error {
                    message: "oops".to_string(),
                },
            ),
        ]);

        assert_eq!(core.attention_count(), 2);
        assert_eq!(core.agent_count(), 3);
    }

    #[test]
    fn test_agents_needing_attention() {
        let core = make_core_with_agents(vec![
            test_agent("main:0.0", AgentStatus::Idle),
            test_agent(
                "main:0.1",
                AgentStatus::AwaitingApproval {
                    approval_type: crate::agents::ApprovalType::FileEdit,
                    details: String::new(),
                },
            ),
        ]);

        let attention = core.agents_needing_attention();
        assert_eq!(attention.len(), 1);
        assert_eq!(attention[0].id, "main:0.1");
    }

    #[test]
    fn test_get_preview() {
        let mut agent = test_agent("main:0.0", AgentStatus::Idle);
        agent.last_content_ansi = "\x1b[32mHello\x1b[0m".to_string();
        agent.last_content = "Hello".to_string();

        let core = make_core_with_agents(vec![agent]);

        let preview = core.get_preview("main:0.0").unwrap();
        assert!(preview.contains("Hello"));

        let content = core.get_content("main:0.0").unwrap();
        assert_eq!(content, "Hello");
    }

    #[test]
    fn test_list_teams_empty() {
        let core = TmaiCoreBuilder::new(Settings::default()).build();
        assert!(core.list_teams().is_empty());
    }

    #[test]
    fn test_is_running() {
        let core = TmaiCoreBuilder::new(Settings::default()).build();
        assert!(core.is_running());
    }

    #[test]
    fn test_get_transcript_no_registry() {
        // Without transcript registry, returns empty vec
        let core = make_core_with_agents(vec![test_agent("main:0.0", AgentStatus::Idle)]);
        let records = core.get_transcript("main:0.0").unwrap();
        assert!(records.is_empty());
    }

    #[test]
    fn test_get_transcript_agent_not_found() {
        let core = TmaiCoreBuilder::new(Settings::default()).build();
        let result = core.get_transcript("nonexistent");
        assert!(matches!(result, Err(ApiError::AgentNotFound { .. })));
    }

    #[test]
    fn test_get_transcript_with_registry() {
        use crate::transcript::types::TranscriptRecord;
        use crate::transcript::watcher::new_transcript_registry;

        let registry = new_transcript_registry();
        // Insert test records
        {
            let mut reg = registry.write();
            let mut state = crate::transcript::TranscriptState::new(
                "/tmp/test.jsonl".to_string(),
                "sess1".to_string(),
                "main:0.0".to_string(),
            );
            state.push_records(vec![
                TranscriptRecord::User {
                    text: "Hello".to_string(),
                },
                TranscriptRecord::AssistantText {
                    text: "Hi there".to_string(),
                },
            ]);
            reg.insert("main:0.0".to_string(), state);
        }

        let app_state = AppState::shared();
        {
            let mut s = app_state.write();
            s.update_agents(vec![test_agent("main:0.0", AgentStatus::Idle)]);
        }

        let core = TmaiCoreBuilder::new(Settings::default())
            .with_state(app_state)
            .with_transcript_registry(registry)
            .build();

        let records = core.get_transcript("main:0.0").unwrap();
        assert_eq!(records.len(), 2);
    }
}