use anyhow::{bail, Context, Result};
use regex::Regex;
use std::fs;
use std::path::{Path, PathBuf};
use tokio::process::Command;
use tracing::{debug, warn};
use crate::worktree::config::{WorktreeConfig, WorktreeMode};
use crate::worktree::status::WorktreeInfo;
#[derive(Debug, Clone)]
pub struct CreateOptions {
pub task_id: String,
pub base_branch: Option<String>,
pub force: bool,
pub custom_path: Option<PathBuf>,
}
#[derive(Debug, Clone)]
pub struct RemoveOptions {
pub target: String,
pub force: bool,
pub delete_branch: bool,
}
impl Default for CreateOptions {
fn default() -> Self {
Self {
task_id: String::new(),
base_branch: None,
force: false,
custom_path: None,
}
}
}
impl Default for RemoveOptions {
fn default() -> Self {
Self {
target: String::new(),
force: false,
delete_branch: false,
}
}
}
#[derive(Debug, Clone)]
pub enum WorktreeOperation {
Create(CreateOptions),
Remove(RemoveOptions),
List,
Status(String),
}
#[derive(Clone)]
pub struct WorktreeOperations {
repo_root: PathBuf,
config: WorktreeConfig,
repo_name: Option<String>,
}
impl WorktreeOperations {
pub fn new(repo_root: PathBuf, config: WorktreeConfig) -> Self {
let repo_name = repo_root
.file_name()
.and_then(|n| n.to_str())
.map(|s| s.to_string());
Self {
repo_root,
config,
repo_name,
}
}
pub fn new_with_repo_name(
repo_root: PathBuf,
config: WorktreeConfig,
repo_name: String,
) -> Self {
Self {
repo_root,
config,
repo_name: Some(repo_name),
}
}
pub async fn create_worktree(&self, options: CreateOptions) -> Result<WorktreeInfo> {
let sanitized_task_id = sanitize_branch_name(&options.task_id)?;
let branch_name = format!("{}{}", self.config.prefix, sanitized_task_id);
validate_branch_name(&branch_name)?;
let worktree_path = match options.custom_path {
Some(custom) => custom,
None => self.calculate_worktree_path(&sanitized_task_id)?,
};
self.ensure_base_directory_exists().await?;
if self.config.auto_gitignore {
self.update_gitignore().await?;
}
let branch_exists = self.branch_exists(&branch_name).await?;
if branch_exists && !options.force {
bail!(
"Branch '{}' already exists. Use --force to recreate.",
branch_name
);
}
let result = if branch_exists && options.force {
if let Ok(existing_path) = self.find_worktree_path(&branch_name).await {
warn!("Removing existing worktree at: {}", existing_path.display());
self.execute_git_command(&[
"worktree",
"remove",
"--force",
&existing_path.to_string_lossy(),
])
.await?;
}
self.execute_git_command(&["branch", "-D", &branch_name])
.await
.ok(); self.create_branch_and_worktree(
&branch_name,
&worktree_path,
options.base_branch.as_deref(),
)
.await?
} else {
self.create_branch_and_worktree(
&branch_name,
&worktree_path,
options.base_branch.as_deref(),
)
.await?
};
debug!(
"Created worktree: {} -> {}",
branch_name,
worktree_path.display()
);
Ok(WorktreeInfo {
path: worktree_path,
branch: branch_name,
head: result.head,
task_id: Some(options.task_id), status: Default::default(), age: std::time::Duration::from_secs(0),
is_detached: false,
})
}
pub async fn remove_worktree(&self, options: RemoveOptions) -> Result<()> {
let worktree_info = self.resolve_worktree_target(&options.target).await?;
let worktree_path = worktree_info.path;
if !worktree_path.exists() {
bail!("Worktree path does not exist: {}", worktree_path.display());
}
if !self.is_valid_worktree(&worktree_path).await? {
bail!(
"Path is not a valid git worktree: {}",
worktree_path.display()
);
}
let branch_name_for_deletion = if options.delete_branch {
Some(worktree_info.branch.clone())
} else {
None
};
let mut args = vec!["worktree", "remove"];
if options.force {
args.push("--force");
}
let path_str = worktree_path.to_string_lossy();
args.push(&path_str);
self.execute_git_command(&args).await?;
if let Some(branch_name) = branch_name_for_deletion {
self.execute_git_command(&["branch", "-D", &branch_name])
.await?;
debug!("Deleted branch: {}", branch_name);
}
debug!("Removed worktree: {}", worktree_path.display());
Ok(())
}
pub async fn list_worktrees(&self) -> Result<Vec<WorktreeInfo>> {
let output = self
.execute_git_command(&["worktree", "list", "--porcelain"])
.await?;
self.parse_worktree_list(&output).await
}
pub async fn find_git_root(&self) -> Result<PathBuf> {
let output = self
.execute_git_command(&["rev-parse", "--show-toplevel"])
.await?;
Ok(PathBuf::from(output.trim()))
}
pub fn get_config(&self) -> &WorktreeConfig {
&self.config
}
pub async fn find_worktree_by_task_id(&self, task_id: &str) -> Result<Option<WorktreeInfo>> {
let worktrees = self.list_worktrees().await?;
Ok(worktrees
.into_iter()
.find(|w| w.task_id.as_ref().map_or(false, |id| id == task_id)))
}
pub async fn resolve_worktree_target(&self, target: &str) -> Result<WorktreeInfo> {
if let Some(worktree) = self.find_worktree_by_task_id(target).await? {
return Ok(worktree);
}
if let Ok(path) = PathBuf::from(target).canonicalize() {
let worktrees = self.list_worktrees().await?;
if let Some(worktree) = worktrees.into_iter().find(|w| w.path == path) {
return Ok(worktree);
}
}
let worktrees = self.list_worktrees().await?;
if let Some(worktree) = worktrees.into_iter().find(|w| w.branch == target) {
return Ok(worktree);
}
bail!("Worktree not found: {}", target)
}
async fn create_branch_and_worktree(
&self,
branch_name: &str,
worktree_path: &Path,
base_branch: Option<&str>,
) -> Result<CreateResult> {
let base = base_branch.unwrap_or("HEAD");
let _output = self
.execute_git_command(&[
"worktree",
"add",
"-b",
branch_name,
&worktree_path.to_string_lossy(),
base,
])
.await?;
let head = self
.execute_git_command(&["rev-parse", "HEAD"])
.await?
.trim()
.to_string();
Ok(CreateResult { head })
}
fn calculate_worktree_path(&self, task_id: &str) -> Result<PathBuf> {
match self.config.mode {
WorktreeMode::Local => {
let base_path = if self.config.base_dir.is_absolute() {
self.config.base_dir.clone()
} else {
self.repo_root.join(&self.config.base_dir)
};
let path_segments: Vec<&str> = task_id.split('/').collect();
let mut worktree_path = base_path;
for segment in path_segments {
worktree_path = worktree_path.join(segment);
}
let timestamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)?
.as_secs();
let final_name = format!(
"{}__{:x}",
worktree_path
.file_name()
.and_then(|n| n.to_str())
.unwrap_or("worktree"),
timestamp
);
Ok(worktree_path
.parent()
.unwrap_or(&worktree_path)
.join(final_name))
}
WorktreeMode::Global => {
let base_path = if self.config.base_dir.is_absolute() {
self.config.base_dir.clone()
} else {
if let Some(home) = dirs::home_dir() {
home.join(".toolprint")
.join("vibe-workspace")
.join("worktrees")
} else {
std::env::temp_dir().join("vibe-worktrees")
}
};
let repo_name = self.repo_name.as_ref().ok_or_else(|| {
anyhow::anyhow!("Repository name required for global worktree mode")
})?;
let timestamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)?
.as_secs();
let safe_task_id = task_id.replace('/', "-");
let worktree_name = format!("{}__{:x}", safe_task_id, timestamp);
Ok(base_path.join(repo_name).join(worktree_name))
}
}
}
async fn ensure_base_directory_exists(&self) -> Result<()> {
let base_path = match self.config.mode {
WorktreeMode::Local => {
if self.config.base_dir.is_absolute() {
self.config.base_dir.clone()
} else {
self.repo_root.join(&self.config.base_dir)
}
}
WorktreeMode::Global => {
let base = if self.config.base_dir.is_absolute() {
self.config.base_dir.clone()
} else {
if let Some(home) = dirs::home_dir() {
home.join(".toolprint")
.join("vibe-workspace")
.join("worktrees")
} else {
std::env::temp_dir().join("vibe-worktrees")
}
};
if let Some(repo_name) = &self.repo_name {
base.join(repo_name)
} else {
base
}
}
};
if !base_path.exists() {
fs::create_dir_all(&base_path).with_context(|| {
format!("Failed to create base directory: {}", base_path.display())
})?;
}
Ok(())
}
async fn update_gitignore(&self) -> Result<()> {
if self.config.mode == WorktreeMode::Global {
return Ok(()); }
let base_path = if self.config.base_dir.is_absolute() {
return Ok(()); } else {
self.config.base_dir.clone()
};
let gitignore_path = self.repo_root.join(".gitignore");
let ignore_pattern = format!("{}/", base_path.display());
if gitignore_path.exists() {
let content = fs::read_to_string(&gitignore_path)?;
if content
.lines()
.any(|line| line.trim() == ignore_pattern.trim())
{
return Ok(()); }
}
let mut content = if gitignore_path.exists() {
fs::read_to_string(&gitignore_path)?
} else {
String::new()
};
if !content.is_empty() && !content.ends_with('\n') {
content.push('\n');
}
content.push_str(&format!(
"# Vibe worktree directories\n{}\n",
ignore_pattern
));
fs::write(&gitignore_path, content).with_context(|| {
format!(
"Failed to update .gitignore at: {}",
gitignore_path.display()
)
})?;
debug!("Updated .gitignore with pattern: {}", ignore_pattern);
Ok(())
}
async fn branch_exists(&self, branch_name: &str) -> Result<bool> {
let result = self
.execute_git_command(&[
"show-ref",
"--verify",
"--quiet",
&format!("refs/heads/{}", branch_name),
])
.await;
Ok(result.is_ok())
}
async fn find_worktree_path(&self, branch_name: &str) -> Result<PathBuf> {
let worktrees = self.list_worktrees().await?;
for worktree in worktrees {
if worktree.branch == branch_name {
return Ok(worktree.path);
}
}
bail!("No worktree found for branch: {}", branch_name);
}
async fn is_valid_worktree(&self, path: &Path) -> Result<bool> {
if !path.exists() {
return Ok(false);
}
let result = Command::new("git")
.args(&["rev-parse", "--show-toplevel"])
.current_dir(path)
.output()
.await?;
Ok(result.status.success())
}
async fn get_worktree_branch(&self, worktree_path: &Path) -> Result<String> {
let output = Command::new("git")
.args(&["rev-parse", "--abbrev-ref", "HEAD"])
.current_dir(worktree_path)
.output()
.await?;
if !output.status.success() {
bail!(
"Failed to get branch name for worktree: {}",
worktree_path.display()
);
}
Ok(String::from_utf8_lossy(&output.stdout).trim().to_string())
}
async fn parse_worktree_list(&self, output: &str) -> Result<Vec<WorktreeInfo>> {
let mut worktrees = Vec::new();
let lines: Vec<&str> = output.lines().collect();
let mut i = 0;
while i < lines.len() {
let line = lines[i];
if line.starts_with("worktree ") {
let path = PathBuf::from(line.strip_prefix("worktree ").unwrap());
let mut branch = String::new();
let mut head = String::new();
let mut is_detached = false;
i += 1;
while i < lines.len() && !lines[i].starts_with("worktree ") {
let info_line = lines[i];
if info_line.starts_with("HEAD ") {
head = info_line.strip_prefix("HEAD ").unwrap().to_string();
} else if info_line.starts_with("branch ") {
let branch_ref = info_line.strip_prefix("branch ").unwrap();
branch = branch_ref
.strip_prefix("refs/heads/")
.unwrap_or(branch_ref)
.to_string();
} else if info_line == "detached" {
is_detached = true;
branch = "(detached)".to_string();
} else if info_line == "bare" {
branch = "(bare)".to_string();
}
i += 1;
}
let age = if let Ok(metadata) = fs::metadata(&path) {
if let Ok(created) = metadata.created() {
std::time::SystemTime::now()
.duration_since(created)
.unwrap_or_default()
} else {
std::time::Duration::from_secs(0)
}
} else {
std::time::Duration::from_secs(0)
};
let task_id = if branch.starts_with(&self.config.prefix) {
Some(
branch
.strip_prefix(&self.config.prefix)
.unwrap_or(&branch)
.to_string(),
)
} else {
None
};
worktrees.push(WorktreeInfo {
path,
branch,
head,
task_id,
status: Default::default(),
age,
is_detached,
});
} else {
i += 1;
}
}
Ok(worktrees)
}
async fn execute_git_command(&self, args: &[&str]) -> Result<String> {
let output = Command::new("git")
.args(args)
.current_dir(&self.repo_root)
.output()
.await
.with_context(|| format!("Failed to execute git command: git {}", args.join(" ")))?;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
bail!(
"Git command failed: git {}\nError: {}",
args.join(" "),
stderr
);
}
Ok(String::from_utf8_lossy(&output.stdout).to_string())
}
}
#[derive(Debug)]
struct CreateResult {
head: String,
}
pub fn validate_branch_name(branch_name: &str) -> Result<()> {
if branch_name.is_empty() {
bail!("Branch name cannot be empty");
}
let dangerous_chars = [
'$', '`', '(', ')', '{', '}', '|', '&', ';', '<', '>', '\n', '\r', '\0', '"', '\'', '\\',
];
if branch_name.chars().any(|c| dangerous_chars.contains(&c)) {
bail!("Branch name contains invalid characters");
}
if branch_name.starts_with('.') || branch_name.ends_with('.') {
bail!("Branch name cannot start or end with a dot");
}
if branch_name.starts_with('/') || branch_name.ends_with('/') {
bail!("Branch name cannot start or end with a slash");
}
if branch_name.contains("..") {
bail!("Branch name cannot contain consecutive dots");
}
if branch_name.contains("@{") {
bail!("Branch name cannot contain '@{{' sequence");
}
if branch_name.len() > 255 {
bail!("Branch name too long (max 255 characters)");
}
Ok(())
}
pub fn sanitize_branch_name(name: &str) -> Result<String> {
if name.is_empty() {
bail!("Task ID cannot be empty");
}
let re = Regex::new(r"[^a-zA-Z0-9\-_/]")?;
let sanitized = re.replace_all(name, "-").to_string();
let re = Regex::new(r"-+")?;
let sanitized = re.replace_all(&sanitized, "-").to_string();
let sanitized = sanitized.trim_matches('-').trim_matches('/');
if sanitized.is_empty() {
bail!(
"Task ID '{}' cannot be sanitized to a valid branch name",
name
);
}
Ok(sanitized.to_string())
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
use tokio;
async fn setup_test_repo() -> Result<(TempDir, PathBuf)> {
let temp_dir = TempDir::new()?;
let repo_path = temp_dir.path().to_path_buf();
Command::new("git")
.args(&["init"])
.current_dir(&repo_path)
.status()
.await?;
Command::new("git")
.args(&["config", "user.email", "test@example.com"])
.current_dir(&repo_path)
.status()
.await?;
Command::new("git")
.args(&["config", "user.name", "Test User"])
.current_dir(&repo_path)
.status()
.await?;
Command::new("git")
.args(&["commit", "--allow-empty", "-m", "Initial commit"])
.current_dir(&repo_path)
.status()
.await?;
Ok((temp_dir, repo_path))
}
#[tokio::test]
async fn test_create_worktree() -> Result<()> {
let (_temp_dir, repo_path) = setup_test_repo().await?;
let config = WorktreeConfig::default();
let ops = WorktreeOperations::new(repo_path, config);
let options = CreateOptions {
task_id: "test-feature".to_string(),
base_branch: None,
force: false,
custom_path: None,
};
let worktree_info = ops.create_worktree(options).await?;
assert!(worktree_info.path.exists());
assert!(worktree_info.branch.starts_with("vibe-ws/"));
assert!(worktree_info.branch.contains("test-feature"));
Ok(())
}
#[tokio::test]
async fn test_list_worktrees() -> Result<()> {
let (_temp_dir, repo_path) = setup_test_repo().await?;
let config = WorktreeConfig::default();
let ops = WorktreeOperations::new(repo_path, config);
let worktrees = ops.list_worktrees().await?;
assert!(!worktrees.is_empty());
Ok(())
}
#[tokio::test]
async fn test_remove_worktree() -> Result<()> {
let (_temp_dir, repo_path) = setup_test_repo().await?;
let config = WorktreeConfig::default();
let ops = WorktreeOperations::new(repo_path, config);
let create_options = CreateOptions {
task_id: "test-remove".to_string(),
base_branch: None,
force: false,
custom_path: None,
};
let worktree_info = ops.create_worktree(create_options).await?;
assert!(worktree_info.path.exists());
let remove_options = RemoveOptions {
target: worktree_info.branch.clone(),
force: false,
delete_branch: true,
};
ops.remove_worktree(remove_options).await?;
assert!(!worktree_info.path.exists());
Ok(())
}
#[tokio::test]
async fn test_path_with_slashes() -> Result<()> {
let (_temp_dir, repo_path) = setup_test_repo().await?;
let config = WorktreeConfig::default();
let ops = WorktreeOperations::new(repo_path, config);
let options = CreateOptions {
task_id: "feat/new-ui".to_string(),
base_branch: None,
force: false,
custom_path: None,
};
let worktree_info = ops.create_worktree(options).await?;
assert!(worktree_info.path.exists());
assert!(worktree_info.path.to_string_lossy().contains("feat"));
Ok(())
}
#[test]
fn test_validate_branch_name() {
assert!(validate_branch_name("feature/new-ui").is_ok());
assert!(validate_branch_name("vibe-ws/task-123").is_ok());
assert!(validate_branch_name("main").is_ok());
assert!(validate_branch_name("").is_err());
assert!(validate_branch_name(".hidden").is_err());
assert!(validate_branch_name("branch.").is_err());
assert!(validate_branch_name("/branch").is_err());
assert!(validate_branch_name("branch/").is_err());
assert!(validate_branch_name("branch..name").is_err());
assert!(validate_branch_name("branch@{upstream}").is_err());
assert!(validate_branch_name("branch$injection").is_err());
assert!(validate_branch_name("branch`command`").is_err());
}
#[test]
fn test_sanitize_branch_name() {
assert_eq!(sanitize_branch_name("Task 123").unwrap(), "Task-123");
assert_eq!(sanitize_branch_name("feat/new-ui").unwrap(), "feat/new-ui");
assert_eq!(
sanitize_branch_name("Fix: issue #456").unwrap(),
"Fix-issue-456"
);
assert_eq!(
sanitize_branch_name("task with spaces").unwrap(),
"task-with-spaces"
);
assert_eq!(
sanitize_branch_name("task---dashes").unwrap(),
"task-dashes"
);
assert!(sanitize_branch_name("").is_err());
assert!(sanitize_branch_name("!!!").is_err());
assert!(sanitize_branch_name("---").is_err());
}
}