Feature: Improved support for standard tmux sessions

.gitignore

@@ -1,3 +1,4 @@
 /target
 /result
 .direnv
+.DS_Store

src/cli.rs

@@ -9,7 +9,7 @@ use crate::{
     execute_command, get_single_selection,
     marks::{marks_command, MarksCommand},
     picker::Preview,
-    session::{create_sessions, SessionContainer},
+    session::{create_all_sessions, create_repo_sessions, SessionContainer},
     tmux::Tmux,
     Result, TmsError,
 };
@@ -314,7 +314,7 @@ fn switch_command(config: Config, tmux: &Tmux) -> Result<()> {
 
     let mut sessions: Vec<String> = sessions.into_iter().map(|s| s.0.to_string()).collect();
     if let Some(true) = config.switch_filter_unknown {
-        let configured = create_sessions(&config)?;
+        let configured = create_repo_sessions(&config)?;
 
         sessions = sessions
             .into_iter()
@@ -791,7 +791,7 @@ fn bookmark_command(args: &BookmarkCommand, mut config: Config) -> Result<()> {
 }
 
 fn open_session_command(args: &OpenSessionCommand, config: Config, tmux: &Tmux) -> Result<()> {
-    let sessions = create_sessions(&config)?;
+    let sessions = create_all_sessions(&config, &tmux)?;
 
     if let Some(session) = sessions.find_session(&args.session) {
         session.switch_to(tmux, &config)?;
@@ -804,7 +804,7 @@ fn open_session_command(args: &OpenSessionCommand, config: Config, tmux: &Tmux)
 fn open_session_completion_candidates() -> Vec<CompletionCandidate> {
     Config::new()
         .change_context(TmsError::ConfigError)
-        .and_then(|config| create_sessions(&config))
+        .and_then(|config| create_repo_sessions(&config))
         .map(|sessions| {
             sessions
                 .list()

src/main.rs

@@ -9,7 +9,7 @@ use tms::{
     error::{Result, Suggestion},
     get_single_selection,
     picker::Preview,
-    session::{create_sessions, SessionContainer},
+    session::{create_all_sessions, SessionContainer},
     tmux::Tmux,
 };
 
@@ -46,17 +46,18 @@ fn main() -> Result<()> {
         SubCommandGiven::No(config) => config, // continue
     };
 
-    let sessions = create_sessions(&config)?;
-    let session_strings = sessions.list();
+    let all_sessions = create_all_sessions(&config, &tmux)?;
+    let all_session_strings = all_sessions.list();
 
-    let selected_str =
-        if let Some(str) = get_single_selection(&session_strings, Preview::None, &config, &tmux)? {
-            str
-        } else {
-            return Ok(());
-        };
+    let selected_str = if let Some(str) =
+        get_single_selection(&all_session_strings, Preview::None, &config, &tmux)?
+    {
+        str
+    } else {
+        return Ok(());
+    };
 
-    if let Some(session) = sessions.find_session(&selected_str) {
+    if let Some(session) = all_sessions.find_session(&selected_str) {
         session.switch_to(&tmux, &config)?;
     }
 

src/session.rs

@@ -23,6 +23,7 @@ pub struct Session {
 pub enum SessionType {
     Git(Repository),
     Bookmark(PathBuf),
+    Standard(PathBuf),
 }
 
 impl Session {
@@ -35,13 +36,15 @@ impl Session {
             SessionType::Git(repo) if repo.is_bare() => repo.path(),
             SessionType::Git(repo) => repo.path().parent().unwrap(),
             SessionType::Bookmark(path) => path,
+            SessionType::Standard(path) => path,
         }
     }
 
     pub fn switch_to(&self, tmux: &Tmux, config: &Config) -> Result<()> {
         match &self.session_type {
             SessionType::Git(repo) => self.switch_to_repo_session(repo, tmux, config),
             SessionType::Bookmark(path) => self.switch_to_bookmark_session(tmux, path, config),
+            SessionType::Standard(path) => self.switch_to_standard_session(tmux, path, config),
         }
     }
 
@@ -85,6 +88,23 @@ impl Session {
 
         Ok(())
     }
+
+    fn switch_to_standard_session(
+        &self,
+        tmux: &Tmux,
+        path: &PathBuf,
+        config: &Config,
+    ) -> Result<()> {
+        let session_name = self.name.to_string();
+
+        if !tmux.session_exists(&session_name) {
+            tmux.new_session(Some(&session_name), path.to_str());
+            tmux.run_session_create_script(path, &session_name, config)?;
+        }
+
+        tmux.switch_to_session(&session_name);
+        Ok(())
+    }
 }
 
 pub trait SessionContainer {
@@ -110,7 +130,7 @@ impl SessionContainer for HashMap<String, Session> {
     }
 }
 
-pub fn create_sessions(config: &Config) -> Result<impl SessionContainer> {
+pub fn create_repo_sessions(config: &Config) -> Result<impl SessionContainer> {
     let mut sessions = find_repos(config)?;
     sessions = append_bookmarks(config, sessions)?;
 
@@ -119,7 +139,22 @@ pub fn create_sessions(config: &Config) -> Result<impl SessionContainer> {
     Ok(sessions)
 }
 
-fn generate_session_container(
+pub fn create_all_sessions(config: &Config, tmux: &Tmux) -> Result<impl SessionContainer> {
+    let repo_sessions = find_repos(config)?;
+    let tmux_sessions = tmux.find_tmux_sessions()?;
+
+    // If session already exists through tmux, dont recommend as a new repo_session
+    let repo_sessions_filtered: HashMap<String, Vec<Session>> = repo_sessions
+        .into_iter()
+        .filter(|(k, _v)| !tmux_sessions.contains_key(k))
+        .collect();
+
+    let all_sessions = merge_session_maps(repo_sessions_filtered, tmux_sessions);
+
+    generate_session_container(all_sessions, config)
+}
+
+pub fn generate_session_container(
     mut sessions: HashMap<String, Vec<Session>>,
     config: &Config,
 ) -> Result<impl SessionContainer> {
@@ -215,7 +250,7 @@ fn deduplicate_sessions(duplicate_sessions: &mut Vec<Session>) -> Vec<Session> {
     deduplicated
 }
 
-fn append_bookmarks(
+pub fn append_bookmarks(
     config: &Config,
     mut sessions: HashMap<String, Vec<Session>>,
 ) -> Result<HashMap<String, Vec<Session>>> {
@@ -237,10 +272,28 @@ fn append_bookmarks(
     Ok(sessions)
 }
 
+fn merge_session_maps(
+    mut s1: HashMap<String, Vec<Session>>,
+    mut s2: HashMap<String, Vec<Session>>,
+) -> HashMap<String, Vec<Session>> {
+    let mut ret: HashMap<String, Vec<Session>> = HashMap::new();
+
+    for (key, mut new_sessions) in s1.drain() {
+        ret.entry(key)
+            .or_insert_with(Vec::new)
+            .append(&mut new_sessions);
+    }
+    for (key, mut new_sessions) in s2.drain() {
+        ret.entry(key)
+            .or_insert_with(Vec::new)
+            .append(&mut new_sessions);
+    }
+    ret
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
-
     #[test]
     fn verify_session_name_deduplication() {
         let mut test_sessions = vec![
@@ -264,4 +317,70 @@ mod tests {
         assert_eq!(deduplicated[1].name, "to/proj2/test");
         assert_eq!(deduplicated[2].name, "to/proj1/test");
     }
+
+    #[test]
+    fn test_merge_session_maps_non_overlapping() {
+        let mut map1 = HashMap::new();
+        let mut map2 = HashMap::new();
+
+        let session_a = Session::new(
+            "Session A".to_string(),
+            SessionType::Standard(PathBuf::from("/path/to/a")),
+        );
+        let session_b = Session::new(
+            "Session B".to_string(),
+            SessionType::Standard(PathBuf::from("/path/to/b")),
+        );
+
+        map1.insert("key1".to_string(), vec![session_a]);
+        map2.insert("key2".to_string(), vec![session_b]);
+
+        let merged = merge_session_maps(map1, map2);
+
+        // Expect both keys to exist independently.
+        assert_eq!(merged.len(), 2);
+        assert!(merged.contains_key("key1"));
+        assert!(merged.contains_key("key2"));
+        assert_eq!(merged["key1"].len(), 1);
+        assert_eq!(merged["key2"].len(), 1);
+    }
+
+    #[test]
+    fn test_merge_session_maps_overlapping() {
+        let mut map1 = HashMap::new();
+        let mut map2 = HashMap::new();
+
+        let session_a = Session::new(
+            "Session A".to_string(),
+            SessionType::Standard(PathBuf::from("/path/to/a")),
+        );
+        let session_b = Session::new(
+            "Session B".to_string(),
+            SessionType::Standard(PathBuf::from("/path/to/b")),
+        );
+
+        // Both maps have the same key "shared_key"
+        map1.insert("shared_key".to_string(), vec![session_a]);
+        map2.insert("shared_key".to_string(), vec![session_b]);
+
+        let merged = merge_session_maps(map1, map2);
+
+        // Expect one key "shared_key" with both sessions, map1's session first.
+        assert_eq!(merged.len(), 1);
+        let sessions = merged.get("shared_key").unwrap();
+        assert_eq!(sessions.len(), 2);
+        assert_eq!(sessions[0].name, "Session A");
+        assert_eq!(sessions[1].name, "Session B");
+    }
+
+    #[test]
+    fn test_merge_session_maps_empty() {
+        let map1: HashMap<String, Vec<Session>> = HashMap::new();
+        let map2: HashMap<String, Vec<Session>> = HashMap::new();
+
+        let merged = merge_session_maps(map1, map2);
+
+        // Expect an empty map.
+        assert!(merged.is_empty());
+    }
 }

src/tmux.rs

@@ -1,4 +1,10 @@
-use std::{env, os::unix::process::CommandExt, path::Path, process};
+use std::{
+    collections::HashMap,
+    env,
+    os::unix::process::CommandExt,
+    path::{Path, PathBuf},
+    process,
+};
 
 use error_stack::ResultExt;
 use git2::Repository;
@@ -7,6 +13,7 @@ use crate::{
     configs::Config,
     dirty_paths::DirtyUtf8Path,
     error::{Result, TmsError},
+    session::{Session, SessionType},
 };
 
 #[derive(Clone)]
@@ -109,18 +116,18 @@ impl Tmux {
         self.replace_with_tmux_command(&args)
     }
 
-    pub fn switch_to_session(&self, repo_short_name: &str) {
+    pub fn switch_to_session(&self, session_name: &str) {
         if !is_in_tmux_session() {
-            self.attach_session(Some(repo_short_name), None);
+            self.attach_session(Some(session_name), None);
         } else {
-            let result = self.switch_client(repo_short_name);
+            let result = self.switch_client(session_name);
             if !result.status.success() {
-                self.attach_session(Some(repo_short_name), None);
+                self.attach_session(Some(session_name), None);
             }
         }
     }
 
-    pub fn session_exists(&self, repo_short_name: &str) -> bool {
+    pub fn session_exists(&self, session_name: &str) -> bool {
         // Get the tmux sessions
         let sessions = self.list_sessions("'#S'");
 
@@ -130,7 +137,7 @@ impl Tmux {
             // tmux will return the output with extra ' and \n characters
             cleaned_line.retain(|char| char != '\'' && char != '\n');
 
-            cleaned_line == repo_short_name
+            cleaned_line == session_name
         })
     }
 
@@ -284,6 +291,33 @@ impl Tmux {
         }
         Ok(())
     }
+
+    pub fn find_tmux_sessions(&self) -> Result<HashMap<String, Vec<Session>>> {
+        let raw = self
+            .list_sessions("'#{session_name},#{pane_start_path}'")
+            .replace('\'', "")
+            .replace("\n\n", "\n");
+
+        let sessions = raw
+            .trim()
+            .lines()
+            .filter_map(|line| line.split_once(','))
+            .map(|(name, path)| {
+                // Trim each part to remove any extraneous whitespace.
+                let session_name = name.trim().to_string();
+                let session_path = PathBuf::from(path.trim());
+                Session::new(session_name, SessionType::Standard(session_path))
+            });
+
+        let grouped = sessions.fold(HashMap::new(), |mut acc, session| {
+            acc.entry(session.name.to_string())
+                .or_insert_with(Vec::new)
+                .push(session);
+            acc
+        });
+
+        Ok(grouped)
+    }
 }
 
 fn is_in_tmux_session() -> bool {