Skip to main content

tmai_core/api/
queries.rs

1//! Read-only query methods on [`TmaiCore`].
2//!
3//! Every method acquires a read lock internally, converts to owned snapshots,
4//! and releases the lock before returning. Callers never hold a lock.
5
6use super::core::TmaiCore;
7use super::types::{AgentDefinitionInfo, AgentSnapshot, ApiError, TeamSummary, TeamTaskInfo};
8
9impl TmaiCore {
10    // =========================================================
11    // Agent queries
12    // =========================================================
13
14    /// List all monitored agents as owned snapshots, in current display order.
15    pub fn list_agents(&self) -> Vec<AgentSnapshot> {
16        let state = self.state().read();
17        let defs = &state.agent_definitions;
18        state
19            .agent_order
20            .iter()
21            .filter_map(|id| state.agents.get(id))
22            .map(|a| {
23                let mut snap = AgentSnapshot::from_agent(a);
24                snap.agent_definition = Self::match_agent_definition(a, defs);
25                snap
26            })
27            .collect()
28    }
29
30    /// Get a single agent snapshot by target ID.
31    pub fn get_agent(&self, target: &str) -> Result<AgentSnapshot, ApiError> {
32        let state = self.state().read();
33        let defs = &state.agent_definitions;
34        state
35            .agents
36            .get(target)
37            .map(|a| {
38                let mut snap = AgentSnapshot::from_agent(a);
39                snap.agent_definition = Self::match_agent_definition(a, defs);
40                snap
41            })
42            .ok_or_else(|| ApiError::AgentNotFound {
43                target: target.to_string(),
44            })
45    }
46
47    /// Get the currently selected agent snapshot.
48    pub fn selected_agent(&self) -> Result<AgentSnapshot, ApiError> {
49        let state = self.state().read();
50        let defs = &state.agent_definitions;
51        state
52            .selected_agent()
53            .map(|agent| {
54                let mut snapshot = AgentSnapshot::from_agent(agent);
55                snapshot.agent_definition = Self::match_agent_definition(agent, defs);
56                snapshot
57            })
58            .ok_or(ApiError::NoSelection)
59    }
60
61    /// Get the number of agents that need user attention.
62    pub fn attention_count(&self) -> usize {
63        let state = self.state().read();
64        state.attention_count()
65    }
66
67    /// Get the total number of monitored agents.
68    pub fn agent_count(&self) -> usize {
69        let state = self.state().read();
70        state.agents.len()
71    }
72
73    /// List agents that need attention (awaiting approval or error).
74    pub fn agents_needing_attention(&self) -> Vec<AgentSnapshot> {
75        let state = self.state().read();
76        state
77            .agent_order
78            .iter()
79            .filter_map(|id| state.agents.get(id))
80            .filter(|a| a.status.needs_attention())
81            .map(AgentSnapshot::from_agent)
82            .collect()
83    }
84
85    // =========================================================
86    // Preview
87    // =========================================================
88
89    /// Get the ANSI preview content for an agent.
90    pub fn get_preview(&self, target: &str) -> Result<String, ApiError> {
91        let state = self.state().read();
92        state
93            .agents
94            .get(target)
95            .map(|a| a.last_content_ansi.clone())
96            .ok_or_else(|| ApiError::AgentNotFound {
97                target: target.to_string(),
98            })
99    }
100
101    /// Get the plain-text content for an agent.
102    pub fn get_content(&self, target: &str) -> Result<String, ApiError> {
103        let state = self.state().read();
104        state
105            .agents
106            .get(target)
107            .map(|a| a.last_content.clone())
108            .ok_or_else(|| ApiError::AgentNotFound {
109                target: target.to_string(),
110            })
111    }
112
113    // =========================================================
114    // Team queries
115    // =========================================================
116
117    /// List all known teams as owned summaries.
118    pub fn list_teams(&self) -> Vec<TeamSummary> {
119        let state = self.state().read();
120        let mut teams: Vec<TeamSummary> = state
121            .teams
122            .values()
123            .map(TeamSummary::from_snapshot)
124            .collect();
125        teams.sort_by(|a, b| a.name.cmp(&b.name));
126        teams
127    }
128
129    /// Get a single team summary by name.
130    pub fn get_team(&self, name: &str) -> Result<TeamSummary, ApiError> {
131        let state = self.state().read();
132        state
133            .teams
134            .get(name)
135            .map(TeamSummary::from_snapshot)
136            .ok_or_else(|| ApiError::TeamNotFound {
137                name: name.to_string(),
138            })
139    }
140
141    /// Get tasks for a team.
142    pub fn get_team_tasks(&self, name: &str) -> Result<Vec<TeamTaskInfo>, ApiError> {
143        let state = self.state().read();
144        state
145            .teams
146            .get(name)
147            .map(|ts| ts.tasks.iter().map(TeamTaskInfo::from_task).collect())
148            .ok_or_else(|| ApiError::TeamNotFound {
149                name: name.to_string(),
150            })
151    }
152
153    // =========================================================
154    // Miscellaneous queries
155    // =========================================================
156
157    /// Match an agent to its definition by configured agent_type or member name.
158    fn match_agent_definition(
159        agent: &crate::agents::MonitoredAgent,
160        defs: &[crate::teams::AgentDefinition],
161    ) -> Option<AgentDefinitionInfo> {
162        if defs.is_empty() {
163            return None;
164        }
165        if let Some(ref team_info) = agent.team_info {
166            // 1) Try configured agent_type (explicit mapping from team config)
167            if let Some(ref agent_type) = team_info.agent_type {
168                if let Some(def) = defs.iter().find(|d| d.name == *agent_type) {
169                    return Some(AgentDefinitionInfo::from_definition(def));
170                }
171            }
172            // 2) Fallback: try member_name as agent definition name
173            if let Some(def) = defs.iter().find(|d| d.name == team_info.member_name) {
174                return Some(AgentDefinitionInfo::from_definition(def));
175            }
176        }
177        None
178    }
179
180    /// Check if the application is still running.
181    pub fn is_running(&self) -> bool {
182        let state = self.state().read();
183        state.running
184    }
185
186    /// Get the last poll timestamp.
187    pub fn last_poll(&self) -> Option<chrono::DateTime<chrono::Utc>> {
188        let state = self.state().read();
189        state.last_poll
190    }
191
192    /// Get known working directories from current agents.
193    pub fn known_directories(&self) -> Vec<String> {
194        let state = self.state().read();
195        state.get_known_directories()
196    }
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202    use crate::agents::{AgentStatus, AgentType, MonitoredAgent};
203    use crate::api::builder::TmaiCoreBuilder;
204    use crate::config::Settings;
205    use crate::state::AppState;
206
207    fn make_core_with_agents(agents: Vec<MonitoredAgent>) -> TmaiCore {
208        let state = AppState::shared();
209        {
210            let mut s = state.write();
211            s.update_agents(agents);
212        }
213        TmaiCoreBuilder::new(Settings::default())
214            .with_state(state)
215            .build()
216    }
217
218    fn test_agent(id: &str, status: AgentStatus) -> MonitoredAgent {
219        let mut agent = MonitoredAgent::new(
220            id.to_string(),
221            AgentType::ClaudeCode,
222            "Title".to_string(),
223            "/home/user".to_string(),
224            100,
225            "main".to_string(),
226            "win".to_string(),
227            0,
228            0,
229        );
230        agent.status = status;
231        agent
232    }
233
234    #[test]
235    fn test_list_agents_empty() {
236        let core = TmaiCoreBuilder::new(Settings::default()).build();
237        assert!(core.list_agents().is_empty());
238    }
239
240    #[test]
241    fn test_list_agents() {
242        let core = make_core_with_agents(vec![
243            test_agent("main:0.0", AgentStatus::Idle),
244            test_agent(
245                "main:0.1",
246                AgentStatus::Processing {
247                    activity: "Bash".to_string(),
248                },
249            ),
250        ]);
251
252        let agents = core.list_agents();
253        assert_eq!(agents.len(), 2);
254    }
255
256    #[test]
257    fn test_get_agent_found() {
258        let core = make_core_with_agents(vec![test_agent("main:0.0", AgentStatus::Idle)]);
259
260        let result = core.get_agent("main:0.0");
261        assert!(result.is_ok());
262        assert_eq!(result.unwrap().id, "main:0.0");
263    }
264
265    #[test]
266    fn test_get_agent_not_found() {
267        let core = TmaiCoreBuilder::new(Settings::default()).build();
268        let result = core.get_agent("nonexistent");
269        assert!(matches!(result, Err(ApiError::AgentNotFound { .. })));
270    }
271
272    #[test]
273    fn test_attention_count() {
274        let core = make_core_with_agents(vec![
275            test_agent("main:0.0", AgentStatus::Idle),
276            test_agent(
277                "main:0.1",
278                AgentStatus::AwaitingApproval {
279                    approval_type: crate::agents::ApprovalType::ShellCommand,
280                    details: "rm -rf".to_string(),
281                },
282            ),
283            test_agent(
284                "main:0.2",
285                AgentStatus::Error {
286                    message: "oops".to_string(),
287                },
288            ),
289        ]);
290
291        assert_eq!(core.attention_count(), 2);
292        assert_eq!(core.agent_count(), 3);
293    }
294
295    #[test]
296    fn test_agents_needing_attention() {
297        let core = make_core_with_agents(vec![
298            test_agent("main:0.0", AgentStatus::Idle),
299            test_agent(
300                "main:0.1",
301                AgentStatus::AwaitingApproval {
302                    approval_type: crate::agents::ApprovalType::FileEdit,
303                    details: String::new(),
304                },
305            ),
306        ]);
307
308        let attention = core.agents_needing_attention();
309        assert_eq!(attention.len(), 1);
310        assert_eq!(attention[0].id, "main:0.1");
311    }
312
313    #[test]
314    fn test_get_preview() {
315        let mut agent = test_agent("main:0.0", AgentStatus::Idle);
316        agent.last_content_ansi = "\x1b[32mHello\x1b[0m".to_string();
317        agent.last_content = "Hello".to_string();
318
319        let core = make_core_with_agents(vec![agent]);
320
321        let preview = core.get_preview("main:0.0").unwrap();
322        assert!(preview.contains("Hello"));
323
324        let content = core.get_content("main:0.0").unwrap();
325        assert_eq!(content, "Hello");
326    }
327
328    #[test]
329    fn test_list_teams_empty() {
330        let core = TmaiCoreBuilder::new(Settings::default()).build();
331        assert!(core.list_teams().is_empty());
332    }
333
334    #[test]
335    fn test_is_running() {
336        let core = TmaiCoreBuilder::new(Settings::default()).build();
337        assert!(core.is_running());
338    }
339}