feat: worktrees at agent session level with random suffix

- Branch name format: kandev/{task_id}_{random8chars}
- Worktree path: ~/.kandev/worktrees/{task_id}_{random}/
- Multiple worktrees per task allowed (one per agent session)
- WorktreeID stored in agent instance metadata for session resumption
- Pass WorktreeID in CreateRequest to reuse existing worktree
- OnTaskDeleted removes ALL worktrees for the task
- Added GetByID, GetAllByTaskID, RemoveByID methods
This commit is contained in:
jcfs
2026-01-12 16:12:34 +00:00
parent 3961afc29e
commit 0b071782c0
6 changed files with 209 additions and 87 deletions
@@ -303,6 +303,7 @@ func (m *Manager) Launch(ctx context.Context, req *LaunchRequest) (*AgentInstanc
// 3. Handle worktree creation/reuse if enabled
workspacePath := req.WorkspacePath
var mainRepoGitDir string // Path to main repo's .git directory for container mount
var worktreeID string // Store worktree ID for session resumption
if req.UseWorktree && m.worktreeMgr != nil && req.RepositoryPath != "" {
wt, err := m.getOrCreateWorktree(ctx, req)
if err != nil {
@@ -312,6 +313,7 @@ func (m *Manager) Launch(ctx context.Context, req *LaunchRequest) (*AgentInstanc
// Fall back to direct mount if worktree creation fails
} else {
workspacePath = wt.Path
worktreeID = wt.ID
// Git worktrees reference the main repo's .git directory via a .git file
// The worktree metadata in .git/worktrees/{name} contains a 'commondir' file
// that points back to the main .git directory (usually '../..')
@@ -319,6 +321,7 @@ func (m *Manager) Launch(ctx context.Context, req *LaunchRequest) (*AgentInstanc
mainRepoGitDir = filepath.Join(req.RepositoryPath, ".git")
m.logger.Info("using worktree for agent",
zap.String("task_id", req.TaskID),
zap.String("worktree_id", worktreeID),
zap.String("worktree_path", workspacePath),
zap.String("main_repo_git_dir", mainRepoGitDir),
zap.String("branch", wt.Branch))
@@ -358,6 +361,14 @@ func (m *Manager) Launch(ctx context.Context, req *LaunchRequest) (*AgentInstanc
// 9. Track the instance
now := time.Now()
instanceMetadata := req.Metadata
if instanceMetadata == nil {
instanceMetadata = make(map[string]interface{})
}
// Store worktree_id in metadata for session resumption
if worktreeID != "" {
instanceMetadata["worktree_id"] = worktreeID
}
instance := &AgentInstance{
ID: instanceID,
TaskID: req.TaskID,
@@ -367,7 +378,7 @@ func (m *Manager) Launch(ctx context.Context, req *LaunchRequest) (*AgentInstanc
Status: v1.AgentStatusRunning,
StartedAt: now,
Progress: 0,
Metadata: req.Metadata,
Metadata: instanceMetadata,
agentctl: agentctlClient,
}
@@ -392,36 +403,25 @@ func (m *Manager) Launch(ctx context.Context, req *LaunchRequest) (*AgentInstanc
return instance, nil
}
// getOrCreateWorktree creates a new worktree or reuses an existing one for the task.
// This enables task resumption - the same worktree is used across multiple agent runs.
// getOrCreateWorktree creates a new worktree or reuses an existing one for session resumption.
// If worktree_id is in metadata, it tries to reuse that specific worktree.
// Otherwise, creates a new worktree with a unique random suffix.
func (m *Manager) getOrCreateWorktree(ctx context.Context, req *LaunchRequest) (*worktree.Worktree, error) {
// First, check if a worktree already exists for this task (session resumption)
existing, err := m.worktreeMgr.GetByTaskID(ctx, req.TaskID)
if err != nil && err != worktree.ErrWorktreeNotFound {
return nil, fmt.Errorf("failed to check existing worktree: %w", err)
}
if existing != nil {
// Validate the existing worktree is still valid
if m.worktreeMgr.IsValid(existing.Path) {
m.logger.Info("reusing existing worktree for task resumption",
zap.String("task_id", req.TaskID),
zap.String("worktree_path", existing.Path),
zap.String("branch", existing.Branch))
return existing, nil
// Check if we have a worktree_id in metadata for session resumption
var worktreeID string
if req.Metadata != nil {
if id, ok := req.Metadata["worktree_id"].(string); ok && id != "" {
worktreeID = id
}
// Worktree exists in DB but is invalid - will be recreated by Create()
m.logger.Warn("existing worktree is invalid, will recreate",
zap.String("task_id", req.TaskID),
zap.String("worktree_path", existing.Path))
}
// Create a new worktree
// Create request with optional WorktreeID for resumption
createReq := worktree.CreateRequest{
TaskID: req.TaskID,
RepositoryID: req.RepositoryID,
RepositoryPath: req.RepositoryPath,
BaseBranch: req.BaseBranch,
WorktreeID: worktreeID, // If set, will try to reuse this worktree
}
wt, err := m.worktreeMgr.Create(ctx, createReq)
@@ -429,10 +429,19 @@ func (m *Manager) getOrCreateWorktree(ctx context.Context, req *LaunchRequest) (
return nil, fmt.Errorf("failed to create worktree: %w", err)
}
m.logger.Info("created new worktree for task",
zap.String("task_id", req.TaskID),
zap.String("worktree_path", wt.Path),
zap.String("branch", wt.Branch))
if worktreeID != "" && wt.ID == worktreeID {
m.logger.Info("reusing existing worktree for session resumption",
zap.String("task_id", req.TaskID),
zap.String("worktree_id", wt.ID),
zap.String("worktree_path", wt.Path),
zap.String("branch", wt.Branch))
} else {
m.logger.Info("created new worktree for task",
zap.String("task_id", req.TaskID),
zap.String("worktree_id", wt.ID),
zap.String("worktree_path", wt.Path),
zap.String("branch", wt.Branch))
}
return wt, nil
}
@@ -62,17 +62,18 @@ func (c *Config) ExpandedBasePath() (string, error) {
return path, nil
}
// WorktreePath returns the full path for a worktree given a task ID.
func (c *Config) WorktreePath(taskID string) (string, error) {
// WorktreePath returns the full path for a worktree given a unique worktree ID.
func (c *Config) WorktreePath(worktreeID string) (string, error) {
basePath, err := c.ExpandedBasePath()
if err != nil {
return "", err
}
return filepath.Join(basePath, taskID), nil
return filepath.Join(basePath, worktreeID), nil
}
// BranchName returns the branch name for a given task ID.
func (c *Config) BranchName(taskID string) string {
return c.BranchPrefix + taskID
// BranchName returns the branch name for a given task ID and suffix.
// Format: {prefix}{taskID}_{suffix} e.g. kandev/abc123_xy7z
func (c *Config) BranchName(taskID, suffix string) string {
return c.BranchPrefix + taskID + "_" + suffix
}
@@ -30,5 +30,8 @@ var (
// ErrGitCommandFailed is returned when a git command fails to execute.
ErrGitCommandFailed = errors.New("git command failed")
// ErrInvalidSession is returned when the session ID is invalid or empty.
ErrInvalidSession = errors.New("invalid or empty session ID")
)
+109 -48
View File
@@ -31,8 +31,12 @@ type Manager struct {
type Store interface {
// CreateWorktree persists a new worktree record.
CreateWorktree(ctx context.Context, wt *Worktree) error
// GetWorktreeByTaskID retrieves a worktree by task ID.
// GetWorktreeByID retrieves a worktree by its unique ID.
GetWorktreeByID(ctx context.Context, id string) (*Worktree, error)
// GetWorktreeByTaskID retrieves the most recent worktree by task ID.
GetWorktreeByTaskID(ctx context.Context, taskID string) (*Worktree, error)
// GetWorktreesByTaskID retrieves all worktrees for a task.
GetWorktreesByTaskID(ctx context.Context, taskID string) ([]*Worktree, error)
// GetWorktreesByRepositoryID retrieves all worktrees for a repository.
GetWorktreesByRepositoryID(ctx context.Context, repoID string) ([]*Worktree, error)
// UpdateWorktree updates an existing worktree record.
@@ -101,26 +105,34 @@ func (m *Manager) GetWorktreeInfo(ctx context.Context, taskID string) (path, bra
return &wt.Path, &wt.Branch
}
// Create creates a new worktree for a task.
// If a worktree already exists for the task, it returns the existing one.
// Create creates a new worktree for a task, or returns an existing one if WorktreeID is provided.
// Each call without WorktreeID creates a new worktree with a unique random suffix.
func (m *Manager) Create(ctx context.Context, req CreateRequest) (*Worktree, error) {
if err := req.Validate(); err != nil {
return nil, err
}
// Check if worktree already exists
existing, err := m.GetByTaskID(ctx, req.TaskID)
if err == nil && existing != nil {
if m.IsValid(existing.Path) {
m.logger.Info("reusing existing worktree",
zap.String("task_id", req.TaskID),
zap.String("path", existing.Path))
return existing, nil
// If WorktreeID is provided, try to reuse that specific worktree (session resumption)
if req.WorktreeID != "" {
existing, err := m.GetByID(ctx, req.WorktreeID)
if err == nil && existing != nil {
if m.IsValid(existing.Path) {
m.logger.Info("reusing existing worktree by ID",
zap.String("worktree_id", req.WorktreeID),
zap.String("task_id", req.TaskID),
zap.String("path", existing.Path))
return existing, nil
}
// Worktree record exists but directory is invalid - recreate
m.logger.Warn("worktree directory invalid, recreating",
zap.String("worktree_id", req.WorktreeID),
zap.String("task_id", req.TaskID))
return m.recreate(ctx, existing, req)
}
// Worktree record exists but directory is invalid - recreate
m.logger.Warn("worktree directory invalid, recreating",
// WorktreeID provided but not found - fall through to create new
m.logger.Warn("worktree ID not found, creating new worktree",
zap.String("worktree_id", req.WorktreeID),
zap.String("task_id", req.TaskID))
return m.recreate(ctx, existing, req)
}
// Check repository is a git repo
@@ -152,15 +164,19 @@ func (m *Manager) Create(ctx context.Context, req CreateRequest) (*Worktree, err
// createWorktree performs the actual git worktree creation.
func (m *Manager) createWorktree(ctx context.Context, req CreateRequest) (*Worktree, error) {
worktreePath, err := m.config.WorktreePath(req.TaskID)
// Generate a unique worktree ID and random suffix
worktreeID := uuid.New().String()
suffix := worktreeID[:8] // Use first 8 chars of UUID as suffix
// Worktree path uses the unique ID: ~/.kandev/worktrees/{task_id}_{suffix}
worktreeDirName := req.TaskID + "_" + suffix
worktreePath, err := m.config.WorktreePath(worktreeDirName)
if err != nil {
return nil, fmt.Errorf("failed to get worktree path: %w", err)
}
branchName := req.BranchName
if branchName == "" {
branchName = m.config.BranchName(req.TaskID)
}
// Branch name: kandev/{task_id}_{suffix}
branchName := m.config.BranchName(req.TaskID, suffix)
// Create worktree with new branch
// git worktree add -b <branch> <path> <base-branch>
@@ -172,31 +188,15 @@ func (m *Manager) createWorktree(ctx context.Context, req CreateRequest) (*Workt
output, err := cmd.CombinedOutput()
if err != nil {
// Check if branch already exists
if strings.Contains(string(output), "already exists") {
// Try without -b flag (use existing branch)
cmd = exec.CommandContext(ctx, "git", "worktree", "add",
worktreePath,
branchName)
cmd.Dir = req.RepositoryPath
output, err = cmd.CombinedOutput()
if err != nil {
m.logger.Error("git worktree add failed",
zap.String("output", string(output)),
zap.Error(err))
return nil, fmt.Errorf("%w: %s", ErrGitCommandFailed, string(output))
}
} else {
m.logger.Error("git worktree add failed",
zap.String("output", string(output)),
zap.Error(err))
return nil, fmt.Errorf("%w: %s", ErrGitCommandFailed, string(output))
}
m.logger.Error("git worktree add failed",
zap.String("output", string(output)),
zap.Error(err))
return nil, fmt.Errorf("%w: %s", ErrGitCommandFailed, string(output))
}
now := time.Now()
wt := &Worktree{
ID: uuid.New().String(),
ID: worktreeID,
TaskID: req.TaskID,
RepositoryID: req.RepositoryID,
RepositoryPath: req.RepositoryPath,
@@ -258,6 +258,30 @@ func (m *Manager) GetByTaskID(ctx context.Context, taskID string) (*Worktree, er
return nil, ErrWorktreeNotFound
}
// GetByID returns a worktree by its unique ID.
func (m *Manager) GetByID(ctx context.Context, worktreeID string) (*Worktree, error) {
if m.store == nil {
return nil, ErrWorktreeNotFound
}
wt, err := m.store.GetWorktreeByID(ctx, worktreeID)
if err != nil {
return nil, err
}
if wt == nil {
return nil, ErrWorktreeNotFound
}
return wt, nil
}
// GetAllByTaskID returns all worktrees for a task.
func (m *Manager) GetAllByTaskID(ctx context.Context, taskID string) ([]*Worktree, error) {
if m.store == nil {
return nil, nil
}
return m.store.GetWorktreesByTaskID(ctx, taskID)
}
// IsValid checks if a worktree directory is valid and usable.
func (m *Manager) IsValid(path string) bool {
// Check directory exists
@@ -281,13 +305,27 @@ func (m *Manager) IsValid(path string) bool {
return true
}
// Remove removes a worktree and optionally its branch.
// Remove removes a worktree by task ID and optionally its branch.
// Deprecated: Use RemoveByID for specific worktree removal.
func (m *Manager) Remove(ctx context.Context, taskID string, removeBranch bool) error {
wt, err := m.GetByTaskID(ctx, taskID)
if err != nil {
return err
}
return m.removeWorktree(ctx, wt, removeBranch)
}
// RemoveByID removes a specific worktree by its ID and optionally its branch.
func (m *Manager) RemoveByID(ctx context.Context, worktreeID string, removeBranch bool) error {
wt, err := m.GetByID(ctx, worktreeID)
if err != nil {
return err
}
return m.removeWorktree(ctx, wt, removeBranch)
}
// removeWorktree performs the actual removal of a worktree.
func (m *Manager) removeWorktree(ctx context.Context, wt *Worktree, removeBranch bool) error {
// Get repository lock
repoLock := m.getRepoLock(wt.RepositoryPath)
repoLock.Lock()
@@ -326,11 +364,12 @@ func (m *Manager) Remove(ctx context.Context, taskID string, removeBranch bool)
// Update cache
m.mu.Lock()
delete(m.worktrees, taskID)
delete(m.worktrees, wt.TaskID)
m.mu.Unlock()
m.logger.Info("removed worktree",
zap.String("task_id", taskID),
zap.String("task_id", wt.TaskID),
zap.String("worktree_id", wt.ID),
zap.String("path", wt.Path),
zap.Bool("branch_removed", removeBranch))
@@ -339,14 +378,36 @@ func (m *Manager) Remove(ctx context.Context, taskID string, removeBranch bool)
// OnTaskDeleted cleans up the worktree when a task is deleted.
// OnTaskDeleted cleans up all worktrees for a task when it is deleted.
func (m *Manager) OnTaskDeleted(ctx context.Context, taskID string) error {
// Try to remove the worktree, ignore if not found
err := m.Remove(ctx, taskID, true)
if err == ErrWorktreeNotFound {
// Get all worktrees for this task
worktrees, err := m.GetAllByTaskID(ctx, taskID)
if err != nil {
return err
}
if len(worktrees) == 0 {
return nil
}
return err
// Remove each worktree and its branch
var lastErr error
for _, wt := range worktrees {
if err := m.removeWorktree(ctx, wt, true); err != nil {
m.logger.Warn("failed to remove worktree on task deletion",
zap.String("task_id", taskID),
zap.String("worktree_id", wt.ID),
zap.Error(err))
lastErr = err
}
}
// Clear from cache
m.mu.Lock()
delete(m.worktrees, taskID)
m.mu.Unlock()
return lastErr
}
// Reconcile syncs worktree state with active tasks on startup.
+51 -3
View File
@@ -29,7 +29,7 @@ func (s *SQLiteStore) initSchema() error {
schema := `
CREATE TABLE IF NOT EXISTS worktrees (
id TEXT PRIMARY KEY,
task_id TEXT NOT NULL UNIQUE,
task_id TEXT NOT NULL,
repository_id TEXT NOT NULL,
repository_path TEXT NOT NULL,
path TEXT NOT NULL,
@@ -76,7 +76,40 @@ func (s *SQLiteStore) CreateWorktree(ctx context.Context, wt *Worktree) error {
return err
}
// GetWorktreeByTaskID retrieves a worktree by task ID.
// GetWorktreeByID retrieves a worktree by its unique ID.
func (s *SQLiteStore) GetWorktreeByID(ctx context.Context, id string) (*Worktree, error) {
wt := &Worktree{}
var mergedAt, deletedAt sql.NullTime
err := s.db.QueryRowContext(ctx, `
SELECT id, task_id, repository_id, repository_path, path, branch, base_branch,
status, created_at, updated_at, merged_at, deleted_at
FROM worktrees WHERE id = ?
`, id).Scan(
&wt.ID, &wt.TaskID, &wt.RepositoryID, &wt.RepositoryPath, &wt.Path,
&wt.Branch, &wt.BaseBranch, &wt.Status, &wt.CreatedAt, &wt.UpdatedAt,
&mergedAt, &deletedAt,
)
if err == sql.ErrNoRows {
return nil, nil // Not found, return nil without error
}
if err != nil {
return nil, err
}
if mergedAt.Valid {
wt.MergedAt = &mergedAt.Time
}
if deletedAt.Valid {
wt.DeletedAt = &deletedAt.Time
}
return wt, nil
}
// GetWorktreeByTaskID retrieves the most recent worktree by task ID.
// Since multiple worktrees can exist per task, this returns the most recently created one.
func (s *SQLiteStore) GetWorktreeByTaskID(ctx context.Context, taskID string) (*Worktree, error) {
wt := &Worktree{}
var mergedAt, deletedAt sql.NullTime
@@ -84,7 +117,7 @@ func (s *SQLiteStore) GetWorktreeByTaskID(ctx context.Context, taskID string) (*
err := s.db.QueryRowContext(ctx, `
SELECT id, task_id, repository_id, repository_path, path, branch, base_branch,
status, created_at, updated_at, merged_at, deleted_at
FROM worktrees WHERE task_id = ?
FROM worktrees WHERE task_id = ? ORDER BY created_at DESC LIMIT 1
`, taskID).Scan(
&wt.ID, &wt.TaskID, &wt.RepositoryID, &wt.RepositoryPath, &wt.Path,
&wt.Branch, &wt.BaseBranch, &wt.Status, &wt.CreatedAt, &wt.UpdatedAt,
@@ -108,6 +141,21 @@ func (s *SQLiteStore) GetWorktreeByTaskID(ctx context.Context, taskID string) (*
return wt, nil
}
// GetWorktreesByTaskID retrieves all worktrees for a task.
func (s *SQLiteStore) GetWorktreesByTaskID(ctx context.Context, taskID string) ([]*Worktree, error) {
rows, err := s.db.QueryContext(ctx, `
SELECT id, task_id, repository_id, repository_path, path, branch, base_branch,
status, created_at, updated_at, merged_at, deleted_at
FROM worktrees WHERE task_id = ? ORDER BY created_at DESC
`, taskID)
if err != nil {
return nil, err
}
defer rows.Close()
return s.scanWorktrees(rows)
}
// GetWorktreesByRepositoryID retrieves all worktrees for a repository.
func (s *SQLiteStore) GetWorktreesByRepositoryID(ctx context.Context, repoID string) ([]*Worktree, error) {
rows, err := s.db.QueryContext(ctx, `
@@ -8,7 +8,7 @@ type Worktree struct {
ID string `json:"id"`
// TaskID is the ID of the task this worktree is associated with.
// This is a 1:1 relationship - each task can have at most one worktree.
// Multiple worktrees can exist for the same task (one per agent session).
TaskID string `json:"task_id"`
// RepositoryID is the ID of the repository this worktree belongs to.
@@ -59,9 +59,9 @@ type CreateRequest struct {
// Typically "main" or "master".
BaseBranch string
// BranchName is the name for the new branch (optional).
// If empty, defaults to "{prefix}{task_id}".
BranchName string
// WorktreeID is the ID of an existing worktree to reuse (optional).
// If provided and valid, the existing worktree is returned instead of creating a new one.
WorktreeID string
}
// Validate validates the create request.