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