1use crate::agents::{AgentStatus, ApprovalType};
7use crate::detectors::get_detector;
8
9use super::core::TmaiCore;
10use super::types::ApiError;
11
12const MAX_TEXT_LENGTH: usize = 1024;
14
15const ALLOWED_KEYS: &[&str] = &[
17 "Enter", "Escape", "Space", "Up", "Down", "Left", "Right", "Tab", "BSpace",
18];
19
20pub fn has_checkbox_format(choices: &[String]) -> bool {
22 choices.iter().any(|c| {
23 let t = c.trim();
24 t.starts_with("[ ]")
25 || t.starts_with("[x]")
26 || t.starts_with("[X]")
27 || t.starts_with("[×]")
28 || t.starts_with("[✔]")
29 })
30}
31
32impl TmaiCore {
33 fn require_command_sender(
39 &self,
40 ) -> Result<&std::sync::Arc<crate::command_sender::CommandSender>, ApiError> {
41 self.command_sender_ref().ok_or(ApiError::NoCommandSender)
42 }
43
44 pub fn approve(&self, target: &str) -> Result<(), ApiError> {
52 let (is_awaiting, agent_type, is_virtual) = {
53 let state = self.state().read();
54 match state.agents.get(target) {
55 Some(a) => (
56 matches!(&a.status, AgentStatus::AwaitingApproval { .. }),
57 a.agent_type.clone(),
58 a.is_virtual,
59 ),
60 None => {
61 return Err(ApiError::AgentNotFound {
62 target: target.to_string(),
63 })
64 }
65 }
66 };
67
68 if is_virtual {
69 return Err(ApiError::VirtualAgent {
70 target: target.to_string(),
71 });
72 }
73
74 if !is_awaiting {
75 return Ok(());
77 }
78
79 let cmd = self.require_command_sender()?;
80 let detector = get_detector(&agent_type);
81 cmd.send_keys(target, detector.approval_keys())?;
82 Ok(())
83 }
84
85 pub fn select_choice(&self, target: &str, choice: usize) -> Result<(), ApiError> {
89 {
91 let state = self.state().read();
92 match state.agents.get(target) {
93 Some(a) if a.is_virtual => {
94 return Err(ApiError::VirtualAgent {
95 target: target.to_string(),
96 });
97 }
98 Some(_) => {}
99 None => {
100 return Err(ApiError::AgentNotFound {
101 target: target.to_string(),
102 });
103 }
104 }
105 }
106
107 let question_info = {
108 let state = self.state().read();
109 state.agents.get(target).and_then(|agent| {
110 if let AgentStatus::AwaitingApproval {
111 approval_type:
112 ApprovalType::UserQuestion {
113 choices,
114 multi_select,
115 cursor_position,
116 },
117 ..
118 } = &agent.status
119 {
120 Some((choices.clone(), *multi_select, *cursor_position))
121 } else {
122 None
123 }
124 })
125 };
126
127 match question_info {
128 Some((choices, multi_select, cursor_pos))
129 if choice >= 1 && choice <= choices.len() + 1 =>
130 {
131 let cmd = self.require_command_sender()?;
132 let cursor = if cursor_pos == 0 { 1 } else { cursor_pos };
133 let steps = choice as i32 - cursor as i32;
134 let key = if steps > 0 { "Down" } else { "Up" };
135 for _ in 0..steps.unsigned_abs() {
136 cmd.send_keys(target, key)?;
137 }
138
139 if !multi_select || has_checkbox_format(&choices) {
141 cmd.send_keys(target, "Enter")?;
142 }
143
144 Ok(())
145 }
146 Some(_) => Err(ApiError::InvalidInput {
147 message: "Invalid choice number".to_string(),
148 }),
149 None => Ok(()),
151 }
152 }
153
154 pub fn submit_selection(
158 &self,
159 target: &str,
160 selected_choices: &[usize],
161 ) -> Result<(), ApiError> {
162 {
164 let state = self.state().read();
165 match state.agents.get(target) {
166 Some(a) if a.is_virtual => {
167 return Err(ApiError::VirtualAgent {
168 target: target.to_string(),
169 });
170 }
171 Some(_) => {}
172 None => {
173 return Err(ApiError::AgentNotFound {
174 target: target.to_string(),
175 });
176 }
177 }
178 }
179
180 let multi_info = {
181 let state = self.state().read();
182 state.agents.get(target).and_then(|agent| {
183 if let AgentStatus::AwaitingApproval {
184 approval_type:
185 ApprovalType::UserQuestion {
186 choices,
187 multi_select: true,
188 cursor_position,
189 },
190 ..
191 } = &agent.status
192 {
193 Some((choices.clone(), *cursor_position))
194 } else {
195 None
196 }
197 })
198 };
199
200 match multi_info {
201 Some((choices, cursor_pos)) => {
202 let cmd = self.require_command_sender()?;
203 let is_checkbox = has_checkbox_format(&choices);
204
205 if is_checkbox && !selected_choices.is_empty() {
206 let mut sorted: Vec<usize> = selected_choices
208 .iter()
209 .copied()
210 .filter(|&c| c >= 1 && c <= choices.len())
211 .collect();
212 if sorted.is_empty() {
213 return Err(ApiError::InvalidInput {
214 message: "No valid choices".to_string(),
215 });
216 }
217 sorted.sort();
218 let mut current_pos = if cursor_pos == 0 { 1 } else { cursor_pos };
219
220 for &choice in &sorted {
221 let steps = choice as i32 - current_pos as i32;
222 let key = if steps > 0 { "Down" } else { "Up" };
223 for _ in 0..steps.unsigned_abs() {
224 cmd.send_keys(target, key)?;
225 }
226 cmd.send_keys(target, "Enter")?;
228 current_pos = choice;
229 }
230 cmd.send_keys(target, "Right")?;
232 cmd.send_keys(target, "Enter")?;
233 } else {
234 let downs_needed = choices.len().saturating_sub(cursor_pos.saturating_sub(1));
236 for _ in 0..downs_needed {
237 cmd.send_keys(target, "Down")?;
238 }
239 cmd.send_keys(target, "Enter")?;
240 }
241 Ok(())
242 }
243 None => Ok(()),
245 }
246 }
247
248 pub async fn send_text(&self, target: &str, text: &str) -> Result<(), ApiError> {
252 if text.chars().count() > MAX_TEXT_LENGTH {
253 return Err(ApiError::InvalidInput {
254 message: format!(
255 "Text exceeds maximum length of {} characters",
256 MAX_TEXT_LENGTH
257 ),
258 });
259 }
260
261 let is_virtual = {
262 let state = self.state().read();
263 match state.agents.get(target) {
264 Some(a) => a.is_virtual,
265 None => {
266 return Err(ApiError::AgentNotFound {
267 target: target.to_string(),
268 })
269 }
270 }
271 };
272
273 if is_virtual {
274 return Err(ApiError::VirtualAgent {
275 target: target.to_string(),
276 });
277 }
278
279 let cmd = self.require_command_sender()?;
280 cmd.send_keys_literal(target, text)?;
281 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
282 cmd.send_keys(target, "Enter")?;
283
284 self.audit_helper()
285 .maybe_emit_input(target, "input_text", "api_input", None);
286
287 Ok(())
288 }
289
290 pub fn send_key(&self, target: &str, key: &str) -> Result<(), ApiError> {
292 if !ALLOWED_KEYS.contains(&key) {
293 return Err(ApiError::InvalidInput {
294 message: "Invalid key name".to_string(),
295 });
296 }
297
298 let is_virtual = {
299 let state = self.state().read();
300 match state.agents.get(target) {
301 Some(a) => a.is_virtual,
302 None => {
303 return Err(ApiError::AgentNotFound {
304 target: target.to_string(),
305 })
306 }
307 }
308 };
309
310 if is_virtual {
311 return Err(ApiError::VirtualAgent {
312 target: target.to_string(),
313 });
314 }
315
316 let cmd = self.require_command_sender()?;
317 cmd.send_keys(target, key)?;
318
319 self.audit_helper()
320 .maybe_emit_input(target, "special_key", "api_input", None);
321
322 Ok(())
323 }
324
325 pub fn focus_pane(&self, target: &str) -> Result<(), ApiError> {
327 {
329 let state = self.state().read();
330 match state.agents.get(target) {
331 Some(a) if a.is_virtual => {
332 return Err(ApiError::VirtualAgent {
333 target: target.to_string(),
334 });
335 }
336 Some(_) => {}
337 None => {
338 return Err(ApiError::AgentNotFound {
339 target: target.to_string(),
340 });
341 }
342 }
343 }
344
345 let cmd = self.require_command_sender()?;
346 cmd.tmux_client().focus_pane(target)?;
347 Ok(())
348 }
349
350 pub fn kill_pane(&self, target: &str) -> Result<(), ApiError> {
352 {
354 let state = self.state().read();
355 match state.agents.get(target) {
356 Some(a) if a.is_virtual => {
357 return Err(ApiError::VirtualAgent {
358 target: target.to_string(),
359 });
360 }
361 Some(_) => {}
362 None => {
363 return Err(ApiError::AgentNotFound {
364 target: target.to_string(),
365 });
366 }
367 }
368 }
369
370 let cmd = self.require_command_sender()?;
371 cmd.tmux_client().kill_pane(target)?;
372 Ok(())
373 }
374}
375
376#[cfg(test)]
377mod tests {
378 use super::*;
379 use crate::agents::{AgentType, MonitoredAgent};
380 use crate::api::builder::TmaiCoreBuilder;
381 use crate::config::Settings;
382 use crate::state::AppState;
383
384 fn make_core_with_agents(agents: Vec<MonitoredAgent>) -> TmaiCore {
385 let state = AppState::shared();
386 {
387 let mut s = state.write();
388 s.update_agents(agents);
389 }
390 TmaiCoreBuilder::new(Settings::default())
391 .with_state(state)
392 .build()
393 }
394
395 fn test_agent(id: &str, status: AgentStatus) -> MonitoredAgent {
396 let mut agent = MonitoredAgent::new(
397 id.to_string(),
398 AgentType::ClaudeCode,
399 "Title".to_string(),
400 "/home/user".to_string(),
401 100,
402 "main".to_string(),
403 "win".to_string(),
404 0,
405 0,
406 );
407 agent.status = status;
408 agent
409 }
410
411 #[test]
412 fn test_has_checkbox_format() {
413 assert!(has_checkbox_format(&[
414 "[ ] Option A".to_string(),
415 "[ ] Option B".to_string(),
416 ]));
417 assert!(has_checkbox_format(&[
418 "[x] Option A".to_string(),
419 "[ ] Option B".to_string(),
420 ]));
421 assert!(has_checkbox_format(&[
422 "[✔] Done".to_string(),
423 "[ ] Not done".to_string(),
424 ]));
425 assert!(!has_checkbox_format(&[
426 "Option A".to_string(),
427 "Option B".to_string(),
428 ]));
429 assert!(!has_checkbox_format(&[]));
430 }
431
432 #[test]
433 fn test_approve_not_found() {
434 let core = TmaiCoreBuilder::new(Settings::default()).build();
435 let result = core.approve("nonexistent");
436 assert!(matches!(result, Err(ApiError::AgentNotFound { .. })));
437 }
438
439 #[test]
440 fn test_approve_virtual_agent() {
441 let mut agent = test_agent(
442 "main:0.0",
443 AgentStatus::AwaitingApproval {
444 approval_type: ApprovalType::FileEdit,
445 details: "edit foo.rs".to_string(),
446 },
447 );
448 agent.is_virtual = true;
449 let core = make_core_with_agents(vec![agent]);
450 let result = core.approve("main:0.0");
451 assert!(matches!(result, Err(ApiError::VirtualAgent { .. })));
452 }
453
454 #[test]
455 fn test_approve_not_awaiting_is_ok() {
456 let agent = test_agent("main:0.0", AgentStatus::Idle);
457 let core = make_core_with_agents(vec![agent]);
458 let result = core.approve("main:0.0");
460 assert!(result.is_ok());
461 }
462
463 #[test]
464 fn test_approve_awaiting_no_command_sender() {
465 let agent = test_agent(
466 "main:0.0",
467 AgentStatus::AwaitingApproval {
468 approval_type: ApprovalType::ShellCommand,
469 details: "rm -rf".to_string(),
470 },
471 );
472 let core = make_core_with_agents(vec![agent]);
473 let result = core.approve("main:0.0");
474 assert!(matches!(result, Err(ApiError::NoCommandSender)));
475 }
476
477 #[test]
478 fn test_send_key_invalid() {
479 let agent = test_agent("main:0.0", AgentStatus::Idle);
480 let core = make_core_with_agents(vec![agent]);
481 let result = core.send_key("main:0.0", "Delete");
482 assert!(matches!(result, Err(ApiError::InvalidInput { .. })));
483 }
484
485 #[test]
486 fn test_send_key_not_found() {
487 let core = TmaiCoreBuilder::new(Settings::default()).build();
488 let result = core.send_key("nonexistent", "Enter");
489 assert!(matches!(result, Err(ApiError::AgentNotFound { .. })));
490 }
491
492 #[test]
493 fn test_send_key_virtual_agent() {
494 let mut agent = test_agent("main:0.0", AgentStatus::Idle);
495 agent.is_virtual = true;
496 let core = make_core_with_agents(vec![agent]);
497 let result = core.send_key("main:0.0", "Enter");
498 assert!(matches!(result, Err(ApiError::VirtualAgent { .. })));
499 }
500
501 #[test]
502 fn test_select_choice_not_in_question() {
503 let agent = test_agent("main:0.0", AgentStatus::Idle);
504 let core = make_core_with_agents(vec![agent]);
505 let result = core.select_choice("main:0.0", 1);
507 assert!(result.is_ok());
508 }
509
510 #[test]
511 fn test_select_choice_not_found() {
512 let core = TmaiCoreBuilder::new(Settings::default()).build();
513 let result = core.select_choice("nonexistent", 1);
514 assert!(matches!(result, Err(ApiError::AgentNotFound { .. })));
515 }
516
517 #[test]
518 fn test_select_choice_virtual_agent() {
519 let mut agent = test_agent("main:0.0", AgentStatus::Idle);
520 agent.is_virtual = true;
521 let core = make_core_with_agents(vec![agent]);
522 let result = core.select_choice("main:0.0", 1);
523 assert!(matches!(result, Err(ApiError::VirtualAgent { .. })));
524 }
525
526 #[test]
527 fn test_select_choice_invalid_number() {
528 let agent = test_agent(
529 "main:0.0",
530 AgentStatus::AwaitingApproval {
531 approval_type: ApprovalType::UserQuestion {
532 choices: vec!["A".to_string(), "B".to_string()],
533 multi_select: false,
534 cursor_position: 1,
535 },
536 details: "Pick one".to_string(),
537 },
538 );
539 let core = make_core_with_agents(vec![agent]);
540 let result = core.select_choice("main:0.0", 0);
542 assert!(matches!(result, Err(ApiError::InvalidInput { .. })));
543 let result = core.select_choice("main:0.0", 4);
545 assert!(matches!(result, Err(ApiError::InvalidInput { .. })));
546 }
547
548 #[tokio::test]
549 async fn test_send_text_too_long() {
550 let agent = test_agent("main:0.0", AgentStatus::Idle);
551 let core = make_core_with_agents(vec![agent]);
552 let long_text = "x".repeat(1025);
553 let result = core.send_text("main:0.0", &long_text).await;
554 assert!(matches!(result, Err(ApiError::InvalidInput { .. })));
555 }
556
557 #[tokio::test]
558 async fn test_send_text_not_found() {
559 let core = TmaiCoreBuilder::new(Settings::default()).build();
560 let result = core.send_text("nonexistent", "hello").await;
561 assert!(matches!(result, Err(ApiError::AgentNotFound { .. })));
562 }
563
564 #[tokio::test]
565 async fn test_send_text_virtual_agent() {
566 let mut agent = test_agent("main:0.0", AgentStatus::Idle);
567 agent.is_virtual = true;
568 let core = make_core_with_agents(vec![agent]);
569 let result = core.send_text("main:0.0", "hello").await;
570 assert!(matches!(result, Err(ApiError::VirtualAgent { .. })));
571 }
572
573 #[tokio::test]
574 async fn test_send_text_at_max_length() {
575 let agent = test_agent("main:0.0", AgentStatus::Idle);
576 let core = make_core_with_agents(vec![agent]);
577 let text = "x".repeat(MAX_TEXT_LENGTH);
579 let result = core.send_text("main:0.0", &text).await;
580 assert!(!matches!(result, Err(ApiError::InvalidInput { .. })));
581 }
582
583 #[test]
584 fn test_focus_pane_not_found() {
585 let core = TmaiCoreBuilder::new(Settings::default()).build();
586 let result = core.focus_pane("nonexistent");
587 assert!(matches!(result, Err(ApiError::AgentNotFound { .. })));
588 }
589
590 #[test]
591 fn test_focus_pane_virtual_agent() {
592 let mut agent = test_agent("main:0.0", AgentStatus::Idle);
593 agent.is_virtual = true;
594 let core = make_core_with_agents(vec![agent]);
595 let result = core.focus_pane("main:0.0");
596 assert!(matches!(result, Err(ApiError::VirtualAgent { .. })));
597 }
598
599 #[test]
600 fn test_kill_pane_not_found() {
601 let core = TmaiCoreBuilder::new(Settings::default()).build();
602 let result = core.kill_pane("nonexistent");
603 assert!(matches!(result, Err(ApiError::AgentNotFound { .. })));
604 }
605
606 #[test]
607 fn test_kill_pane_virtual_agent() {
608 let mut agent = test_agent("main:0.0", AgentStatus::Idle);
609 agent.is_virtual = true;
610 let core = make_core_with_agents(vec![agent]);
611 let result = core.kill_pane("main:0.0");
612 assert!(matches!(result, Err(ApiError::VirtualAgent { .. })));
613 }
614
615 #[test]
616 fn test_submit_selection_not_found() {
617 let core = TmaiCoreBuilder::new(Settings::default()).build();
618 let result = core.submit_selection("nonexistent", &[1]);
619 assert!(matches!(result, Err(ApiError::AgentNotFound { .. })));
620 }
621
622 #[test]
623 fn test_submit_selection_virtual_agent() {
624 let mut agent = test_agent("main:0.0", AgentStatus::Idle);
625 agent.is_virtual = true;
626 let core = make_core_with_agents(vec![agent]);
627 let result = core.submit_selection("main:0.0", &[1]);
628 assert!(matches!(result, Err(ApiError::VirtualAgent { .. })));
629 }
630
631 #[test]
632 fn test_submit_selection_not_in_multiselect() {
633 let agent = test_agent("main:0.0", AgentStatus::Idle);
634 let core = make_core_with_agents(vec![agent]);
635 let result = core.submit_selection("main:0.0", &[1]);
637 assert!(result.is_ok());
638 }
639}