1use crate::cmd::CommandErrorContext;
10use anyhow::Context;
11use std::collections::HashMap;
12use tokio::process::Command;
13
14#[derive(Debug, Clone)]
16pub struct Session {
17 pub name: String,
18 pub task_id: String,
20 pub project: String,
22 pub created_at: Option<u64>, }
24
25#[derive(Clone)]
27pub struct TmuxManager {
28 prefix: String,
30 socket: Option<String>,
32}
33
34impl TmuxManager {
35 pub fn new() -> Self {
36 #[cfg(test)]
37 {
38 Self {
39 prefix: "orch-test-".to_string(),
40 socket: Some(test_socket_name()),
41 }
42 }
43 #[cfg(not(test))]
44 {
45 Self {
46 prefix: "orch-".to_string(),
47 socket: None,
48 }
49 }
50 }
51
52 pub(crate) fn command(&self) -> Command {
54 let mut cmd = Command::new("tmux");
55 if let Some(socket) = &self.socket {
56 cmd.args(["-L", socket]);
57 }
58 cmd
59 }
60
61 pub fn session_name(&self, project: &str, task_id: &str) -> String {
66 let project_name = project
67 .rsplit('/')
68 .next()
69 .unwrap_or(project)
70 .trim_end_matches(".git");
71 let safe_task_id = task_id.replace(':', "-");
72 format!("{}{project_name}-{safe_task_id}", self.prefix)
73 }
74
75 pub async fn create_session(
90 &self,
91 repo: &str,
92 task_id: &str,
93 working_dir: &str,
94 command: &str,
95 env_vars: &[(&str, &str)],
96 ) -> anyhow::Result<String> {
97 let name = self.session_name(repo, task_id);
98
99 let mut cmd = self.command();
100 cmd.args(["new-session", "-d", "-s", &name, "-c", working_dir]);
101
102 cmd.arg("-e").arg("DISABLE_AUTO_UPDATE=true");
104
105 for (key, value) in env_vars {
109 cmd.arg("-e").arg(format!("{}={}", key, value));
111 }
112
113 cmd.arg(command);
114
115 let output = cmd
116 .output_with_context()
117 .await
118 .context("spawning tmux session")?;
119
120 if !output.status.success() {
121 let stderr = String::from_utf8_lossy(&output.stderr);
122 anyhow::bail!("tmux new-session failed: {stderr}");
123 }
124
125 tracing::info!(session = %name, task_id, "created tmux session");
126 Ok(name)
127 }
128
129 pub async fn session_exists(&self, session: &str) -> bool {
131 self.command()
132 .args(["has-session", "-t", session])
133 .output_with_context()
134 .await
135 .map(|o| o.status.success())
136 .unwrap_or(false)
137 }
138
139 pub async fn kill_session(&self, session: &str) -> anyhow::Result<()> {
141 let output = self
142 .command()
143 .args(["kill-session", "-t", session])
144 .output_with_context()
145 .await?;
146
147 if !output.status.success() {
148 let stderr = String::from_utf8_lossy(&output.stderr);
149 if stderr.contains("can't find session")
150 || stderr.contains("no server running")
151 || stderr.contains("error connecting to")
152 {
153 tracing::debug!(session, %stderr, "kill-session: session already gone");
154 } else {
155 tracing::warn!(session, %stderr, "kill-session failed unexpectedly");
156 }
157 }
158 Ok(())
159 }
160
161 pub async fn capture_pane(&self, session: &str, lines: i32) -> anyhow::Result<String> {
167 let start = format!("-{lines}");
168 let output = self
169 .command()
170 .args(["capture-pane", "-t", session, "-p", "-S", &start])
171 .output_with_context()
172 .await?;
173
174 if !output.status.success() {
175 let stderr = String::from_utf8_lossy(&output.stderr);
176 if stderr.contains("can't find pane")
181 || stderr.contains("can't find session")
182 || stderr.contains("no server running")
183 || stderr.contains("error connecting to")
184 {
185 tracing::debug!(
186 session,
187 "capture_pane: session/pane already gone (teardown race)"
188 );
189 return Ok(String::new());
190 }
191 anyhow::bail!("capture-pane failed for {session}: {stderr}");
192 }
193 Ok(String::from_utf8_lossy(&output.stdout).to_string())
194 }
195
196 pub async fn is_session_active(&self, session: &str) -> bool {
199 let output = self
201 .command()
202 .args(["list-panes", "-t", session, "-F", "#{pane_dead}"])
203 .output_with_context()
204 .await;
205
206 match output {
207 Ok(o) if o.status.success() => {
208 let flag = String::from_utf8_lossy(&o.stdout);
209 flag.trim() == "0" }
211 _ => false,
212 }
213 }
214
215 #[cfg(test)]
217 pub async fn session_is_running(&self, session: &str) -> bool {
218 self.session_exists(session).await && self.is_session_active(session).await
219 }
220
221 pub async fn session_blocks_dispatch(&self, session: &str) -> bool {
226 if !self.session_exists(session).await {
227 return false;
228 }
229
230 if self.is_session_active(session).await {
231 return true;
232 }
233
234 tracing::warn!(
235 session,
236 "found stale tmux session with dead pane; cleaning up"
237 );
238 if let Err(err) = self.kill_session(session).await {
239 tracing::debug!(session, error = %err, "failed to clean up stale tmux session");
240 }
241
242 false
243 }
244
245 pub async fn session_blocks_dispatch_from_map(
251 &self,
252 session: &str,
253 session_map: &std::collections::HashMap<String, bool>,
254 ) -> bool {
255 match session_map.get(session).copied() {
256 None => false, Some(true) => true, Some(false) => {
259 tracing::warn!(
261 session,
262 "found stale tmux session with dead pane; cleaning up"
263 );
264 if let Err(err) = self.kill_session(session).await {
265 tracing::debug!(
266 session,
267 error = %err,
268 "failed to clean up stale tmux session"
269 );
270 }
271 false
272 }
273 }
274 }
275
276 pub(crate) fn parse_session_name(&self, name: &str) -> Option<(String, String)> {
290 let after_prefix = name.strip_prefix(&self.prefix)?;
291
292 let (base, is_review) = match after_prefix.strip_suffix("-review") {
294 Some(stripped) => (stripped, true),
295 None => (after_prefix, false),
296 };
297
298 if let Some(idx) = base.rfind("-internal-") {
300 let project = &base[..idx];
301 let task_base = &base[idx + 1..]; let after_internal = &task_base["internal-".len()..];
303 if !project.is_empty()
304 && !after_internal.is_empty()
305 && after_internal.chars().all(|c| c.is_ascii_digit())
306 {
307 let task_id = if is_review {
308 format!("{task_base}-review")
309 } else {
310 task_base.to_string()
311 };
312 return Some((project.to_string(), task_id));
313 }
314 }
315
316 if let Some(idx) = base.rfind('-') {
318 let project = &base[..idx];
319 let digits = &base[idx + 1..];
320 if !project.is_empty()
321 && !digits.is_empty()
322 && digits.chars().all(|c| c.is_ascii_digit())
323 {
324 let task_id = if is_review {
325 format!("{digits}-review")
326 } else {
327 digits.to_string()
328 };
329 return Some((project.to_string(), task_id));
330 }
331 }
332
333 None
334 }
335
336 pub async fn batch_session_active(&self) -> HashMap<String, bool> {
341 let mut map = HashMap::new();
342 let o = match self
343 .command()
344 .args(["list-panes", "-a", "-F", "#{session_name} #{pane_dead}"])
345 .output_with_context()
346 .await
347 {
348 Ok(o) => o,
349 Err(e) => {
350 tracing::warn!(error = %e, "batch_session_active: failed to spawn tmux");
351 return map;
352 }
353 };
354 if !o.status.success() {
355 let stderr = String::from_utf8_lossy(&o.stderr);
356 tracing::warn!(
357 status = ?o.status,
358 stderr = %stderr.trim(),
359 "batch_session_active: tmux list-panes exited non-zero"
360 );
361 return map;
362 }
363 for line in String::from_utf8_lossy(&o.stdout).lines() {
364 let mut parts = line.trim().splitn(2, ' ');
365 if let (Some(session), Some(dead)) = (parts.next(), parts.next()) {
366 map.insert(session.to_string(), dead.trim() == "0");
367 }
368 }
369 map
370 }
371
372 pub async fn list_sessions(&self) -> anyhow::Result<Vec<Session>> {
374 let output = self
376 .command()
377 .args(["list-sessions", "-F", "#{session_name} #{session_created}"])
378 .output_with_context()
379 .await?;
380
381 if !output.status.success() {
382 return Ok(vec![]);
383 }
384
385 let mut sessions = Vec::new();
386 for line in String::from_utf8_lossy(&output.stdout).lines() {
387 let parts: Vec<&str> = line.trim().splitn(2, ' ').collect();
388 if parts.is_empty() {
389 continue;
390 }
391 let name = parts[0].to_string();
392 let created_at = parts.get(1).and_then(|s| s.parse::<u64>().ok());
393
394 if let Some((project, task_id)) = self.parse_session_name(&name) {
395 sessions.push(Session {
396 name,
397 task_id,
398 project,
399 created_at,
400 });
401 }
402 }
403
404 Ok(sessions)
405 }
406
407 pub async fn kill_stale_sessions(&self, before_timestamp: u64) -> anyhow::Result<usize> {
410 let sessions = self.list_sessions().await?;
411 let mut killed = 0;
412
413 for session in sessions {
414 let should_kill = match session.created_at {
415 Some(created) => created < before_timestamp,
416 None => false,
418 };
419
420 if should_kill {
421 tracing::info!(
422 session = %session.name,
423 created_at = session.created_at,
424 "killing stale tmux session from previous run"
425 );
426 if let Err(e) = self.kill_session(&session.name).await {
427 tracing::warn!(session = %session.name, error = %e, "failed to kill stale session");
428 } else {
429 killed += 1;
430 }
431 }
432 }
433
434 if killed > 0 {
435 tracing::info!(killed, "killed stale tmux sessions on startup");
436 }
437
438 Ok(killed)
439 }
440
441 pub async fn wait_for_sessions_dead(
445 &self,
446 sessions: &[String],
447 poll_interval: std::time::Duration,
448 timeout: std::time::Duration,
449 ) -> usize {
450 let start = std::time::Instant::now();
451 let mut remaining: std::collections::HashSet<String> = sessions.iter().cloned().collect();
452
453 while !remaining.is_empty() && start.elapsed() < timeout {
454 let mut to_remove = Vec::new();
455 for session in &remaining {
456 if !self.session_exists(session).await {
457 to_remove.push(session.clone());
458 }
459 }
460 for session in to_remove {
461 remaining.remove(&session);
462 }
463 if !remaining.is_empty() {
464 tokio::time::sleep(poll_interval).await;
465 }
466 }
467
468 let still_alive = remaining.len();
469 if still_alive > 0 {
470 tracing::warn!(
471 still_alive,
472 sessions = ?remaining.iter().collect::<Vec<_>>(),
473 "tmux sessions did not terminate within timeout"
474 );
475 }
476 still_alive
477 }
478
479 pub async fn wait_for_completion(
482 &self,
483 session: &str,
484 poll_interval: std::time::Duration,
485 ) -> anyhow::Result<String> {
486 loop {
487 if !self.session_exists(session).await {
488 return Ok(String::new());
489 }
490
491 if !self.is_session_active(session).await {
492 match self.capture_pane(session, 500).await {
496 Ok(output) => return Ok(output),
497 Err(err) => {
498 tracing::error!(session = %session, error = %err, "failed to capture final pane output");
499 return Ok(String::new());
500 }
501 }
502 }
503
504 tokio::time::sleep(poll_interval).await;
505 }
506 }
507
508 pub async fn snapshot(&self) -> Vec<(Session, bool)> {
516 let sessions = match self.list_sessions().await {
517 Ok(s) => s,
518 Err(err) => {
519 tracing::error!(error = %err, "failed to list tmux sessions — returning empty snapshot");
520 Vec::new()
521 }
522 };
523 let active_map = self.batch_session_active().await;
525 sessions
526 .into_iter()
527 .map(|s| {
528 let active = active_map.get(&s.name).copied().unwrap_or(false);
530 (s, active)
531 })
532 .collect()
533 }
534
535 #[cfg(test)]
543 pub async fn set_env(&self, session: &str, key: &str, value: &str) -> anyhow::Result<()> {
544 let output = self
545 .command()
546 .args(["set-environment", "-t", session, key, value])
547 .output_with_context()
548 .await?;
549
550 if !output.status.success() {
551 let stderr = String::from_utf8_lossy(&output.stderr);
552 anyhow::bail!("tmux set-environment failed for {session}: {stderr}");
553 }
554
555 tracing::debug!(session, key, "set tmux environment variable");
556 Ok(())
557 }
558
559 pub async fn unset_global_env(&self, key: &str) -> anyhow::Result<()> {
563 let output = self
564 .command()
565 .args(["set-environment", "-gu", key])
566 .output_with_context()
567 .await?;
568
569 if !output.status.success() {
570 let stderr = String::from_utf8_lossy(&output.stderr);
571 anyhow::bail!("tmux unset-global-environment failed for {key}: {stderr}");
572 }
573
574 tracing::debug!(key, "unset tmux global environment variable");
575 Ok(())
576 }
577}
578
579#[cfg(test)]
580pub(crate) fn test_session_name() -> String {
581 format!(
582 "orch-test-{}",
583 std::time::SystemTime::now()
584 .duration_since(std::time::UNIX_EPOCH)
585 .map(|d| d.as_millis())
586 .unwrap_or_else(|_| 0)
587 )
588}
589
590#[cfg(test)]
591fn test_socket_name() -> String {
592 format!("orch-test-{}", std::process::id())
593}
594
595#[cfg(test)]
597pub(crate) struct SessionGuard(pub(crate) String);
598
599#[cfg(test)]
600impl Drop for SessionGuard {
601 fn drop(&mut self) {
602 let _ = std::process::Command::new("tmux")
603 .args(["-L", &test_socket_name(), "kill-session", "-t", &self.0])
604 .output();
605 }
606}
607
608#[cfg(test)]
609mod tests {
610 use super::*;
611
612 #[test]
613 fn session_name_internal_task_id_is_sanitized() {
614 let tmux = TmuxManager::new();
615 let name = tmux.session_name("owner/repo", "internal:8");
616 assert!(!name.contains(':'), "colon must be sanitized for tmux");
617 assert_eq!(
618 tmux.parse_session_name(&name),
619 Some(("repo".to_string(), "internal-8".to_string()))
620 );
621 }
622
623 #[test]
624 fn task_id_from_session_name_parses_external_task() {
625 let tmux = TmuxManager::new();
626 let name = tmux.session_name("owner/repo", "1234");
627 assert_eq!(
628 tmux.parse_session_name(&name),
629 Some(("repo".to_string(), "1234".to_string()))
630 );
631 }
632
633 #[test]
634 fn task_id_from_session_name_parses_internal_task() {
635 let tmux = TmuxManager::new();
636 let name = tmux.session_name("owner/orch", "internal:21116");
637 assert_eq!(
638 tmux.parse_session_name(&name),
639 Some(("orch".to_string(), "internal-21116".to_string()))
640 );
641 }
642
643 #[test]
644 fn task_id_from_session_name_parses_internal_task_small_id() {
645 let tmux = TmuxManager::new();
646 let name = tmux.session_name("owner/repo", "internal:8");
647 assert_eq!(
648 tmux.parse_session_name(&name),
649 Some(("repo".to_string(), "internal-8".to_string()))
650 );
651 }
652
653 #[test]
654 fn task_id_from_session_name_returns_none_for_non_orch_session() {
655 let tmux = TmuxManager::new();
656 assert_eq!(tmux.parse_session_name("some-other-session"), None);
657 }
658
659 #[test]
660 fn parse_session_name_ignores_production_prefix_in_tests() {
661 let tmux = TmuxManager::new();
662 assert_eq!(tmux.parse_session_name("orch-repo-1234"), None);
663 }
664
665 #[test]
666 fn parse_session_name_hyphenated_project_external() {
667 let tmux = TmuxManager::new();
668 let name = tmux.session_name("owner/my-repo", "1234");
669 assert_eq!(
670 tmux.parse_session_name(&name),
671 Some(("my-repo".to_string(), "1234".to_string()))
672 );
673 }
674
675 #[test]
676 fn parse_session_name_hyphenated_project_internal() {
677 let tmux = TmuxManager::new();
678 let name = tmux.session_name("owner/my-repo", "internal:8");
679 assert_eq!(
680 tmux.parse_session_name(&name),
681 Some(("my-repo".to_string(), "internal-8".to_string()))
682 );
683 }
684
685 #[test]
686 fn parse_session_name_hyphenated_project_review() {
687 let tmux = TmuxManager::new();
688 let name = tmux.session_name("owner/my-repo", "1234-review");
689 assert_eq!(
690 tmux.parse_session_name(&name),
691 Some(("my-repo".to_string(), "1234-review".to_string()))
692 );
693 }
694
695 #[test]
696 fn parse_session_name_multi_hyphen_project() {
697 let tmux = TmuxManager::new();
698 let name = tmux.session_name("owner/my-cool-repo", "42");
699 assert_eq!(
700 tmux.parse_session_name(&name),
701 Some(("my-cool-repo".to_string(), "42".to_string()))
702 );
703 }
704
705 #[test]
706 fn parse_session_name_roundtrip_hyphenated_project() {
707 let tmux = TmuxManager::new();
708 let name = tmux.session_name("owner/my-repo", "1793");
709 assert!(name.ends_with("my-repo-1793"));
710 assert_eq!(
711 tmux.parse_session_name(&name),
712 Some(("my-repo".to_string(), "1793".to_string()))
713 );
714 }
715
716 #[tokio::test]
719 async fn test_set_env() {
720 let tmux = TmuxManager::new();
721 let session = test_session_name();
722
723 let create_result = tmux
725 .command()
726 .args(["new-session", "-d", "-s", &session, "-c", "/tmp"])
727 .output()
728 .await;
729
730 match create_result {
732 Ok(o) if o.status.success() => {}
733 _ => {
734 eprintln!("Skipping test: tmux not available or failed to create test session");
735 return;
736 }
737 }
738 let _guard = SessionGuard(session.clone());
739
740 let result = tmux.set_env(&session, "TEST_VAR", "test_value").await;
742 assert!(result.is_ok(), "set_env should succeed");
743
744 let check_result = tmux
746 .command()
747 .args(["show-environment", "-t", &session, "TEST_VAR"])
748 .output()
749 .await;
750
751 let output = match check_result {
752 Ok(o) => o,
753 Err(_) => {
754 eprintln!("Skipping test: unable to check tmux environment");
755 return;
756 }
757 };
758 let stdout = String::from_utf8_lossy(&output.stdout);
759 assert!(
760 stdout.contains("TEST_VAR=test_value"),
761 "Expected TEST_VAR=test_value, got: {}",
762 stdout
763 );
764 }
765
766 #[tokio::test]
767 async fn session_blocks_dispatch_keeps_live_session() {
768 let tmux = TmuxManager::new();
769 let session = test_session_name();
770
771 let create_result = tmux
772 .command()
773 .args(["new-session", "-d", "-s", &session, "-c", "/tmp"])
774 .output()
775 .await;
776
777 match create_result {
778 Ok(o) if o.status.success() => {}
779 _ => {
780 eprintln!("Skipping test: tmux not available or failed to create test session");
781 return;
782 }
783 }
784 let _guard = SessionGuard(session.clone());
785
786 assert!(tmux.session_blocks_dispatch(&session).await);
787 assert!(tmux.session_exists(&session).await);
788 }
789
790 #[tokio::test]
791 async fn session_blocks_dispatch_cleans_dead_session() {
792 let tmux = TmuxManager::new();
793 let session = test_session_name();
794
795 let create_result = tmux
796 .command()
797 .args(["new-session", "-d", "-s", &session, "-c", "/tmp"])
798 .output()
799 .await;
800
801 match create_result {
802 Ok(o) if o.status.success() => {}
803 _ => {
804 eprintln!("Skipping test: tmux not available or failed to create test session");
805 return;
806 }
807 }
808 let _guard = SessionGuard(session.clone());
809
810 let set_option_result = tmux
811 .command()
812 .args(["set-option", "-t", &session, "remain-on-exit", "on"])
813 .output()
814 .await;
815 if !matches!(set_option_result, Ok(ref o) if o.status.success()) {
816 eprintln!("Skipping test: unable to set tmux remain-on-exit option");
817 return;
818 }
819
820 let respawn_result = tmux
823 .command()
824 .args(["respawn-pane", "-k", "-t", &session, "false"])
825 .output()
826 .await;
827 if !matches!(respawn_result, Ok(ref o) if o.status.success()) {
828 eprintln!("Skipping test: unable to exit tmux pane");
829 return;
830 }
831
832 for _ in 0..20 {
834 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
835 if !tmux.session_is_running(&session).await {
836 break;
837 }
838 }
839
840 assert!(tmux.session_exists(&session).await);
841 assert!(!tmux.session_is_running(&session).await);
842 assert!(!tmux.session_blocks_dispatch(&session).await);
843 assert!(!tmux.session_exists(&session).await);
844 }
845
846 #[tokio::test]
850 async fn capture_pane_returns_empty_for_gone_session() {
851 let tmux = TmuxManager::new();
852 let result = tmux
854 .capture_pane("orch-nonexistent-session-99999", 100)
855 .await;
856 assert!(
857 result.is_ok(),
858 "capture_pane should return Ok for non-existent session, got: {:?}",
859 result
860 );
861 assert_eq!(
862 result.unwrap(),
863 "",
864 "capture_pane should return empty string for non-existent session"
865 );
866 }
867
868 #[tokio::test]
872 async fn wait_for_completion_returns_empty_for_gone_session() {
873 let tmux = TmuxManager::new();
874 let result = tmux
876 .wait_for_completion(
877 "orch-nonexistent-session-99999",
878 std::time::Duration::from_millis(50),
879 )
880 .await;
881 assert!(
882 result.is_ok(),
883 "wait_for_completion should return Ok for non-existent session, got: {:?}",
884 result
885 );
886 assert_eq!(
887 result.unwrap(),
888 "",
889 "wait_for_completion should return empty string for non-existent session"
890 );
891 }
892
893 #[tokio::test]
896 async fn test_manager_does_not_touch_the_default_server() {
897 let tmux = TmuxManager::new();
898 let name = match tmux
899 .create_session("owner/repo", "isolation", "/tmp", "sleep 30", &[])
900 .await
901 {
902 Ok(n) => n,
903 Err(_) => {
904 eprintln!("Skipping test: tmux not available or failed to create test session");
905 return;
906 }
907 };
908 let _guard = SessionGuard(name.clone());
909
910 assert!(
911 tmux.session_exists(&name).await,
912 "session must exist on the isolated test server"
913 );
914
915 let on_default = tokio::process::Command::new("tmux")
916 .args(["has-session", "-t", &name])
917 .output()
918 .await
919 .map(|o| o.status.success())
920 .unwrap_or(false);
921 assert!(
922 !on_default,
923 "TmuxManager under cfg(test) must not create sessions on the default tmux server"
924 );
925 }
926}