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