1use super::core::TmaiCore;
7use super::types::{AgentDefinitionInfo, AgentSnapshot, ApiError, TeamSummary, TeamTaskInfo};
8
9impl TmaiCore {
10 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 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 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 pub fn attention_count(&self) -> usize {
63 let state = self.state().read();
64 state.attention_count()
65 }
66
67 pub fn agent_count(&self) -> usize {
69 let state = self.state().read();
70 state.agents.len()
71 }
72
73 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 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 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 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 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 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 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 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 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 pub fn is_running(&self) -> bool {
182 let state = self.state().read();
183 state.running
184 }
185
186 pub fn last_poll(&self) -> Option<chrono::DateTime<chrono::Utc>> {
188 let state = self.state().read();
189 state.last_poll
190 }
191
192 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}