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