1use std::collections::BTreeMap;
2use std::ffi::OsString;
3use std::path::{Path, PathBuf};
4
5use crate::git::Git;
6use crate::{Error, Result};
7
8const ROOT_ENV_KEYS: &[&str] = &[
9 "TREEBOOT_ROOT_PATH",
10 "CODEX_SOURCE_TREE_PATH",
11 "CONDUCTOR_ROOT_PATH",
12 "SUPERSET_ROOT_PATH",
13];
14
15pub type Environment = BTreeMap<String, OsString>;
17
18#[derive(Debug, Clone, Default, PartialEq, Eq)]
20pub struct WorktreeOptions {
21 pub cwd: Option<PathBuf>,
23 pub root: Option<PathBuf>,
25}
26
27#[derive(Debug, Clone, PartialEq, Eq)]
29pub struct Worktree {
30 pub root_path: PathBuf,
32 pub worktree_path: PathBuf,
34 pub default_branch: String,
36 pub environment: Environment,
38}
39
40impl Worktree {
41 pub fn discover(options: WorktreeOptions) -> Result<Self> {
49 resolve(&options)
50 }
51}
52
53pub(crate) fn resolve(options: &WorktreeOptions) -> Result<Worktree> {
54 let cwd = options.cwd.as_ref().map_or_else(
55 || std::env::current_dir().map_err(|source| Error::CurrentDir { source }),
56 |path| Ok(path.clone()),
57 )?;
58 let git = Git::new(&cwd);
59 let worktree_path = normalize_existing_path(&git.worktree_path()?)?;
60 let root_path = discover_root_path(options, &cwd, &git)?;
61 let default_branch = discover_default_branch(&git)?;
62 let environment = build_environment(&root_path, &worktree_path, &default_branch);
63
64 Ok(Worktree {
65 root_path,
66 worktree_path,
67 default_branch,
68 environment,
69 })
70}
71
72fn discover_root_path(options: &WorktreeOptions, cwd: &Path, git: &Git) -> Result<PathBuf> {
73 if let Some(path) = &options.root {
74 return normalize_existing_path(&resolve_input_path(cwd, path));
75 }
76
77 for key in ROOT_ENV_KEYS {
78 if let Some(value) = non_empty_env(key) {
79 return normalize_existing_path(&resolve_input_path(cwd, &PathBuf::from(value)));
80 }
81 }
82
83 git.main_worktree_path()?
84 .map(|path| normalize_existing_path(&path))
85 .transpose()?
86 .ok_or(Error::RootPathNotFound)
87}
88
89fn discover_default_branch(git: &Git) -> Result<String> {
90 if let Some(branch) = non_empty_env("CONDUCTOR_DEFAULT_BRANCH") {
91 return Ok(branch.to_string_lossy().into_owned());
92 }
93
94 git.default_branch()
95}
96
97fn build_environment(root_path: &Path, worktree_path: &Path, default_branch: &str) -> Environment {
98 let root = root_path.as_os_str().to_os_string();
99 let worktree = worktree_path.as_os_str().to_os_string();
100 let branch = OsString::from(default_branch);
101
102 let mut env = Environment::new();
103 env.insert("TREEBOOT_ROOT_PATH".to_owned(), root.clone());
104 env.insert("TREEBOOT_WORKTREE_PATH".to_owned(), worktree.clone());
105 env.insert("TREEBOOT_DEFAULT_BRANCH".to_owned(), branch.clone());
106 env.insert("GIT_SOURCE_TREE_PATH".to_owned(), root.clone());
107 env.insert("GIT_WORKTREE_PATH".to_owned(), worktree.clone());
108 env.insert("CODEX_SOURCE_TREE_PATH".to_owned(), root.clone());
109 env.insert("CODEX_WORKTREE_PATH".to_owned(), worktree.clone());
110 env.insert("CONDUCTOR_ROOT_PATH".to_owned(), root.clone());
111 env.insert("CONDUCTOR_WORKSPACE_PATH".to_owned(), worktree);
112 env.insert("CONDUCTOR_DEFAULT_BRANCH".to_owned(), branch);
113 env.insert("SUPERSET_ROOT_PATH".to_owned(), root);
114 env
115}
116
117fn non_empty_env(key: &str) -> Option<OsString> {
118 std::env::var_os(key).filter(|value| !value.is_empty())
119}
120
121pub(crate) fn resolve_worktree_path(worktree_path: &Path, path: &Path) -> PathBuf {
122 if path.is_absolute() {
123 path.to_path_buf()
124 } else {
125 worktree_path.join(path)
126 }
127}
128
129fn resolve_input_path(cwd: &Path, path: &Path) -> PathBuf {
130 if path.is_absolute() {
131 path.to_path_buf()
132 } else {
133 cwd.join(path)
134 }
135}
136
137fn normalize_existing_path(path: &Path) -> Result<PathBuf> {
138 std::fs::canonicalize(path).map_err(|source| Error::NormalizePath {
139 path: path.to_path_buf(),
140 source,
141 })
142}
143
144#[cfg(test)]
145mod tests {
146 use super::*;
147
148 #[test]
149 fn build_environment_should_set_codex_worktree_to_worktree_path() {
150 let root = Path::new("/repo");
151 let worktree = Path::new("/repo-worktree");
152 let env = build_environment(root, worktree, "main");
153
154 assert_eq!(
155 env.get("CODEX_WORKTREE_PATH"),
156 Some(&OsString::from("/repo-worktree"))
157 );
158 }
159
160 #[test]
161 fn resolve_worktree_path_should_join_relative_paths() {
162 let worktree = Path::new("/repo-worktree");
163
164 assert_eq!(
165 resolve_worktree_path(worktree, Path::new(".env")),
166 PathBuf::from("/repo-worktree/.env")
167 );
168 }
169
170 #[test]
171 fn resolve_worktree_path_should_keep_absolute_paths() {
172 let worktree = Path::new("/repo-worktree");
173
174 assert_eq!(
175 resolve_worktree_path(worktree, Path::new("/tmp/.env")),
176 PathBuf::from("/tmp/.env")
177 );
178 }
179}