diff --git a/apps/backend/internal/agent/runtime/lifecycle/env_preparer_worktree_multi_test.go b/apps/backend/internal/agent/runtime/lifecycle/env_preparer_worktree_multi_test.go index 5683df5295..8f8a61f41d 100644 --- a/apps/backend/internal/agent/runtime/lifecycle/env_preparer_worktree_multi_test.go +++ b/apps/backend/internal/agent/runtime/lifecycle/env_preparer_worktree_multi_test.go @@ -116,6 +116,10 @@ func (s *inMemoryWorktreeStore) ListActiveWorktreePaths(_ context.Context) ([]st return paths, nil } +func (s *inMemoryWorktreeStore) CountActiveWorktreeReferences(_ context.Context, _ string, _ []string) (int, error) { + return 0, nil +} + // GetWorktreesBySessionID — MultiRepoStore. func (s *inMemoryWorktreeStore) GetWorktreesBySessionID(_ context.Context, sessionID string) ([]*worktree.Worktree, error) { s.mu.Lock() diff --git a/apps/backend/internal/task/service/service.go b/apps/backend/internal/task/service/service.go index 25dc052686..e9a92e08bc 100644 --- a/apps/backend/internal/task/service/service.go +++ b/apps/backend/internal/task/service/service.go @@ -36,6 +36,11 @@ type WorktreeBatchCleaner interface { CleanupWorktrees(ctx context.Context, worktrees []*worktree.Worktree) error } +type worktreeReferenceGuard interface { + CountActiveWorktreeReferences(ctx context.Context, worktreeID string, excludeSessionIDs []string) (int, error) + ReleaseWorktreeReference(ctx context.Context, wt *worktree.Worktree) error +} + // TaskExecutionStopper stops active task execution (agent session + instance). type TaskExecutionStopper interface { StopTask(ctx context.Context, taskID, reason string, force bool) error diff --git a/apps/backend/internal/task/service/service_tasks.go b/apps/backend/internal/task/service/service_tasks.go index b681c31c40..55d4c8b1cf 100644 --- a/apps/backend/internal/task/service/service_tasks.go +++ b/apps/backend/internal/task/service/service_tasks.go @@ -1763,6 +1763,9 @@ func (s *Service) cleanupDestructiveTaskResources( originalWorktreeCount := len(worktrees) worktrees = s.filterOwnedWorktreesForTaskCleanup(ctx, taskID, sessions, worktrees, envCleanup.env, skipOwnedEnvironment) worktrees = cleanupEligibleWorktrees(worktrees, envCleanup.env, preserveExecutorRows) + var referenceErrs []error + worktrees, referenceErrs = s.filterSharedWorktreesForTaskCleanup(ctx, taskID, sessions, worktrees) + errs = append(errs, referenceErrs...) if len(worktrees) == 0 { if originalWorktreeCount > 0 { s.logger.Debug("no task worktrees eligible for cleanup", @@ -1788,6 +1791,51 @@ func (s *Service) cleanupDestructiveTaskResources( return errs } +func (s *Service) filterSharedWorktreesForTaskCleanup( + ctx context.Context, + taskID string, + sessions []*models.TaskSession, + worktrees []*worktree.Worktree, +) ([]*worktree.Worktree, []error) { + guard, ok := s.worktreeCleanup.(worktreeReferenceGuard) + if !ok || len(worktrees) == 0 { + return worktrees, nil + } + excludedSessionIDs := taskSessionIDs(sessions) + filtered := worktrees[:0] + var errs []error + for _, wt := range worktrees { + count, err := guard.CountActiveWorktreeReferences(ctx, wt.ID, excludedSessionIDs) + if err != nil { + errs = append(errs, fmt.Errorf("count active references for worktree %s: %w", wt.ID, err)) + continue + } + if count == 0 { + filtered = append(filtered, wt) + continue + } + if err := guard.ReleaseWorktreeReference(ctx, wt); err != nil { + errs = append(errs, fmt.Errorf("release shared worktree reference %s: %w", wt.ID, err)) + continue + } + s.logger.Info("preserving worktree still referenced by another active task", + zap.String("task_id", taskID), + zap.String("worktree_id", wt.ID), + zap.Int("active_references", count)) + } + return filtered, errs +} + +func taskSessionIDs(sessions []*models.TaskSession) []string { + ids := make([]string, 0, len(sessions)) + for _, session := range sessions { + if session != nil && session.ID != "" { + ids = append(ids, session.ID) + } + } + return ids +} + func taskEnvironmentID(env *models.TaskEnvironment) string { if env == nil { return "" diff --git a/apps/backend/internal/task/service/service_test.go b/apps/backend/internal/task/service/service_test.go index 89d0e8cf73..ff3d48ab9f 100644 --- a/apps/backend/internal/task/service/service_test.go +++ b/apps/backend/internal/task/service/service_test.go @@ -2014,6 +2014,9 @@ type recordingWorktreeCleanup struct { worktrees []*worktree.Worktree worktreesByTaskID map[string][]*worktree.Worktree cleaned []*worktree.Worktree + referenceCounts map[string]int + released []*worktree.Worktree + excludedSessions []string } func (c *recordingWorktreeCleanup) OnTaskDeleted(context.Context, string) error { @@ -2034,6 +2037,20 @@ func (c *recordingWorktreeCleanup) CleanupWorktrees(_ context.Context, worktrees return nil } +func (c *recordingWorktreeCleanup) CountActiveWorktreeReferences(_ context.Context, worktreeID string, excludeSessionIDs []string) (int, error) { + c.mu.Lock() + defer c.mu.Unlock() + c.excludedSessions = append([]string(nil), excludeSessionIDs...) + return c.referenceCounts[worktreeID], nil +} + +func (c *recordingWorktreeCleanup) ReleaseWorktreeReference(_ context.Context, wt *worktree.Worktree) error { + c.mu.Lock() + defer c.mu.Unlock() + c.released = append(c.released, wt) + return nil +} + func (c *recordingWorktreeCleanup) cleanedIDs() []string { c.mu.Lock() defer c.mu.Unlock() @@ -2046,6 +2063,50 @@ func (c *recordingWorktreeCleanup) cleanedIDs() []string { return ids } +func (c *recordingWorktreeCleanup) releasedIDs() []string { + c.mu.Lock() + defer c.mu.Unlock() + ids := make([]string, 0, len(c.released)) + for _, wt := range c.released { + if wt != nil { + ids = append(ids, wt.ID) + } + } + return ids +} + +func TestService_CleanupDestructiveTaskResourcesPreservesSharedWorktree(t *testing.T) { + svc, _, _ := createTestService(t) + cleanup := &recordingWorktreeCleanup{ + referenceCounts: map[string]int{"worktree-shared": 1}, + } + svc.SetWorktreeCleanup(cleanup) + + session := &models.TaskSession{ID: "session-owner", TaskID: "task-owner"} + wt := &worktree.Worktree{ + ID: "worktree-shared", + TaskID: session.TaskID, + SessionID: session.ID, + } + errs := svc.cleanupDestructiveTaskResources( + context.Background(), session.TaskID, []*models.TaskSession{session}, + []*worktree.Worktree{wt}, taskEnvironmentCleanup{}, nil, + ) + + if len(errs) != 0 { + t.Fatalf("cleanup errors = %v, want none", errs) + } + if got := cleanup.cleanedIDs(); len(got) != 0 { + t.Fatalf("physically cleaned worktrees = %v, want none", got) + } + if got := cleanup.releasedIDs(); len(got) != 1 || got[0] != wt.ID { + t.Fatalf("released worktrees = %v, want [%s]", got, wt.ID) + } + if got := cleanup.excludedSessions; len(got) != 1 || got[0] != session.ID { + t.Fatalf("excluded sessions = %v, want [%s]", got, session.ID) + } +} + func waitForCleanupDone(t *testing.T, svc *Service) { t.Helper() select { diff --git a/apps/backend/internal/worktree/manager.go b/apps/backend/internal/worktree/manager.go index da23b3bbcd..b2bf5b8c57 100644 --- a/apps/backend/internal/worktree/manager.go +++ b/apps/backend/internal/worktree/manager.go @@ -100,6 +100,9 @@ type Store interface { // non-deleted worktree row that has a non-empty path. Used by the // office GC to identify live worktrees that must not be swept. ListActiveWorktreePaths(ctx context.Context) ([]string, error) + // CountActiveWorktreeReferences counts non-terminal session associations + // for a physical worktree, excluding associations owned by the caller. + CountActiveWorktreeReferences(ctx context.Context, worktreeID string, excludeSessionIDs []string) (int, error) } // MultiRepoStore is an optional capability some stores implement to support @@ -176,6 +179,16 @@ func (m *Manager) ListActiveWorktreePaths(ctx context.Context) ([]string, error) return m.store.ListActiveWorktreePaths(ctx) } +// CountActiveWorktreeReferences returns the number of live session +// associations for a physical worktree outside the supplied sessions. +func (m *Manager) CountActiveWorktreeReferences( + ctx context.Context, + worktreeID string, + excludeSessionIDs []string, +) (int, error) { + return m.store.CountActiveWorktreeReferences(ctx, worktreeID, excludeSessionIDs) +} + // IsEnabled returns whether worktree mode is enabled. func (m *Manager) IsEnabled() bool { return m.config.Enabled diff --git a/apps/backend/internal/worktree/manager_cleanup.go b/apps/backend/internal/worktree/manager_cleanup.go index 0e3926458b..5f20366bf5 100644 --- a/apps/backend/internal/worktree/manager_cleanup.go +++ b/apps/backend/internal/worktree/manager_cleanup.go @@ -95,6 +95,20 @@ func (m *Manager) removeWorktree(ctx context.Context, wt *Worktree, removeBranch repoLock.Unlock() m.releaseRepoLock(wt.RepositoryPath) }() + activeReferences, err := m.CountActiveWorktreeReferences(ctx, wt.ID, []string{wt.SessionID}) + if err != nil { + return fmt.Errorf("count active references for worktree %s: %w", wt.ID, err) + } + if activeReferences > 0 { + if err := m.ReleaseWorktreeReference(ctx, wt); err != nil { + return fmt.Errorf("release shared worktree reference %s: %w", wt.ID, err) + } + m.logger.Info("preserved worktree still referenced by another active session", + zap.String("worktree_id", wt.ID), + zap.String("session_id", wt.SessionID), + zap.Int("active_references", activeReferences)) + return nil + } // Execute cleanup script BEFORE removing directory m.runWorktreeCleanupScript(ctx, wt) @@ -129,11 +143,7 @@ func (m *Manager) removeWorktree(ctx context.Context, wt *Worktree, removeBranch // Update store if m.store != nil { - now := time.Now() - wt.Status = StatusDeleted - wt.DeletedAt = &now - wt.UpdatedAt = now - if err := m.store.UpdateWorktree(ctx, wt); err != nil { + if err := m.ReleaseWorktreeReference(ctx, wt); err != nil { // Record may already be deleted by another cleanup path (e.g. task deletion). // This is expected and harmless - only log at debug level. m.logger.Debug("failed to update worktree status (may already be deleted)", @@ -159,6 +169,25 @@ func (m *Manager) removeWorktree(ctx context.Context, wt *Worktree, removeBranch return nil } +// ReleaseWorktreeReference marks one session's association deleted without +// removing the shared directory or branch. +func (m *Manager) ReleaseWorktreeReference(ctx context.Context, wt *Worktree) error { + if wt == nil || wt.SessionID == "" { + return fmt.Errorf("session ID is required to release worktree reference") + } + now := time.Now().UTC() + wt.Status = StatusDeleted + wt.DeletedAt = &now + wt.UpdatedAt = now + if err := m.store.UpdateWorktree(ctx, wt); err != nil && !errors.Is(err, ErrWorktreeNotFound) { + return err + } + m.mu.Lock() + delete(m.worktrees, cacheKey(wt.SessionID, wt.RepositoryID, wt.BranchSlug)) + m.mu.Unlock() + return nil +} + // runWorktreeSetupScript runs the repository's setup script in the freshly // created worktree. Setup script failures are non-fatal: the worktree is kept // and the failure is recorded on wt as a warning (surfaced by the env preparer) diff --git a/apps/backend/internal/worktree/manager_cleanup_shared_reference_test.go b/apps/backend/internal/worktree/manager_cleanup_shared_reference_test.go new file mode 100644 index 0000000000..e345c04fd6 --- /dev/null +++ b/apps/backend/internal/worktree/manager_cleanup_shared_reference_test.go @@ -0,0 +1,168 @@ +package worktree + +import ( + "context" + "os" + "path/filepath" + "testing" + + "github.com/jmoiron/sqlx" + + "github.com/kandev/kandev/internal/db" + "github.com/kandev/kandev/internal/task/models" + "github.com/kandev/kandev/internal/task/repository" +) + +func TestCleanupWorktrees_PreservesSharedActiveReference(t *testing.T) { + mgr, store := newReferenceCleanupTestManager(t) + ctx := context.Background() + seedReferenceCleanupSession(t, store, "task-owner", "session-owner", models.TaskSessionStateRunning) + seedReferenceCleanupSession(t, store, "task-borrower", "session-borrower", models.TaskSessionStateRunning) + + wt := createReferenceCleanupWorktree(t, mgr, "task-owner", "session-owner") + borrowed := *wt + borrowed.TaskID = "task-borrower" + borrowed.SessionID = "session-borrower" + if err := store.CreateWorktree(ctx, &borrowed); err != nil { + t.Fatalf("create borrower worktree reference: %v", err) + } + count, err := store.CountActiveWorktreeReferences(ctx, wt.ID, []string{"session-owner"}) + if err != nil { + t.Fatalf("count active worktree references: %v", err) + } + if count != 1 { + t.Fatalf("active foreign worktree references = %d, want 1", count) + } + + if err := mgr.CleanupWorktrees(ctx, []*Worktree{wt}); err != nil { + t.Fatalf("CleanupWorktrees: %v", err) + } + + if _, err := os.Stat(wt.Path); err != nil { + t.Fatalf("shared worktree path should be preserved: %v", err) + } + assertWorktreeReferenceStatus(t, store, wt.ID, "session-owner", StatusDeleted) + assertWorktreeReferenceStatus(t, store, wt.ID, "session-borrower", StatusActive) +} + +func TestCleanupWorktrees_RemovesLastActiveReference(t *testing.T) { + mgr, store := newReferenceCleanupTestManager(t) + seedReferenceCleanupSession(t, store, "task-owner", "session-owner", models.TaskSessionStateCompleted) + wt := createReferenceCleanupWorktree(t, mgr, "task-owner", "session-owner") + + if err := mgr.CleanupWorktrees(context.Background(), []*Worktree{wt}); err != nil { + t.Fatalf("CleanupWorktrees: %v", err) + } + + if _, err := os.Stat(wt.Path); !os.IsNotExist(err) { + t.Fatalf("last-reference worktree path should be removed, stat error = %v", err) + } + assertWorktreeReferenceStatus(t, store, wt.ID, "session-owner", StatusDeleted) +} + +func TestReleaseWorktreeReference_MissingAssociationIsIdempotent(t *testing.T) { + mgr, store := newReferenceCleanupTestManager(t) + ctx := context.Background() + seedReferenceCleanupSession(t, store, "task-owner", "session-owner", models.TaskSessionStateCompleted) + seedReferenceCleanupSession(t, store, "task-borrower", "session-borrower", models.TaskSessionStateRunning) + wt := createReferenceCleanupWorktree(t, mgr, "task-owner", "session-owner") + borrowed := *wt + borrowed.TaskID = "task-borrower" + borrowed.SessionID = "session-borrower" + if err := store.CreateWorktree(ctx, &borrowed); err != nil { + t.Fatalf("create borrower worktree reference: %v", err) + } + if _, err := store.db.ExecContext(ctx, `DELETE FROM task_session_worktrees WHERE session_id = ?`, wt.SessionID); err != nil { + t.Fatalf("delete owner worktree reference: %v", err) + } + + if err := mgr.ReleaseWorktreeReference(ctx, wt); err != nil { + t.Fatalf("release missing owner reference: %v", err) + } + if _, err := os.Stat(wt.Path); err != nil { + t.Fatalf("shared worktree path should be preserved: %v", err) + } + assertWorktreeReferenceStatus(t, store, wt.ID, borrowed.SessionID, StatusActive) +} + +func newReferenceCleanupTestManager(t *testing.T) (*Manager, *SQLiteStore) { + t.Helper() + dbConn, err := db.OpenSQLite(filepath.Join(t.TempDir(), "cleanup.db")) + if err != nil { + t.Fatalf("open sqlite: %v", err) + } + sqlxDB := sqlx.NewDb(dbConn, "sqlite3") + t.Cleanup(func() { _ = sqlxDB.Close() }) + taskRepo, cleanup, err := repository.Provide(sqlxDB, sqlxDB, nil) + if err != nil { + t.Fatalf("create task repository: %v", err) + } + t.Cleanup(func() { _ = cleanup() }) + store, err := NewSQLiteStore(sqlxDB, sqlxDB) + if err != nil { + t.Fatalf("create worktree store: %v", err) + } + if err := taskRepo.CreateWorkspace(context.Background(), &models.Workspace{ID: "workspace", Name: "Workspace"}); err != nil { + t.Fatalf("create workspace: %v", err) + } + mgr, err := NewManager(newTestConfig(t), store, newTestLogger()) + if err != nil { + t.Fatalf("NewManager: %v", err) + } + return mgr, store +} + +func seedReferenceCleanupSession( + t *testing.T, + store *SQLiteStore, + taskID string, + sessionID string, + state models.TaskSessionState, +) { + t.Helper() + ctx := context.Background() + if _, err := store.db.ExecContext(ctx, ` + INSERT INTO tasks (id, workspace_id, title, priority, created_at, updated_at) + VALUES (?, 'workspace', ?, 'medium', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) + `, taskID, taskID); err != nil { + t.Fatalf("create task %s: %v", taskID, err) + } + if _, err := store.db.ExecContext(ctx, ` + INSERT INTO task_sessions (id, task_id, state, started_at, updated_at) + VALUES (?, ?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) + `, sessionID, taskID, state); err != nil { + t.Fatalf("create session %s: %v", sessionID, err) + } +} + +func createReferenceCleanupWorktree(t *testing.T, mgr *Manager, taskID, sessionID string) *Worktree { + t.Helper() + wt, err := mgr.Create(context.Background(), CreateRequest{ + TaskID: taskID, + SessionID: sessionID, + TaskTitle: "Shared cleanup", + RepositoryID: "repository", + RepositoryPath: initGitRepoWithRemote(t), + BaseBranch: "main", + TaskDirName: taskID, + RepoName: "repository", + }) + if err != nil { + t.Fatalf("create worktree: %v", err) + } + return wt +} + +func assertWorktreeReferenceStatus(t *testing.T, store *SQLiteStore, worktreeID, sessionID, want string) { + t.Helper() + var got string + if err := store.ro.QueryRowContext(context.Background(), ` + SELECT status FROM task_session_worktrees + WHERE worktree_id = ? AND session_id = ? + `, worktreeID, sessionID).Scan(&got); err != nil { + t.Fatalf("load worktree reference status: %v", err) + } + if got != want { + t.Fatalf("worktree reference status = %q, want %q", got, want) + } +} diff --git a/apps/backend/internal/worktree/manager_test.go b/apps/backend/internal/worktree/manager_test.go index d6e678a38f..e96447ef30 100644 --- a/apps/backend/internal/worktree/manager_test.go +++ b/apps/backend/internal/worktree/manager_test.go @@ -113,6 +113,10 @@ func (s *mockStore) ListActiveWorktreePaths(_ context.Context) ([]string, error) return paths, nil } +func (s *mockStore) CountActiveWorktreeReferences(_ context.Context, _ string, _ []string) (int, error) { + return 0, nil +} + // GetWorktreesBySessionID — MultiRepoStore. func (s *mockStore) GetWorktreesBySessionID(_ context.Context, sessionID string) ([]*Worktree, error) { var out []*Worktree diff --git a/apps/backend/internal/worktree/store.go b/apps/backend/internal/worktree/store.go index 4f076216ce..dd924f1998 100644 --- a/apps/backend/internal/worktree/store.go +++ b/apps/backend/internal/worktree/store.go @@ -10,6 +10,7 @@ import ( "github.com/jmoiron/sqlx" "github.com/kandev/kandev/internal/db" + "github.com/kandev/kandev/internal/task/models" ) // SQLiteStore implements Store interface using SQLite. @@ -373,7 +374,7 @@ func (s *SQLiteStore) UpdateWorktree(ctx context.Context, wt *Worktree) error { rows, _ := result.RowsAffected() if rows == 0 { - return fmt.Errorf("worktree not found: %s", wt.ID) + return fmt.Errorf("%w: %s", ErrWorktreeNotFound, wt.ID) } return nil } @@ -452,6 +453,42 @@ func (s *SQLiteStore) ListActiveWorktreePaths(ctx context.Context) ([]string, er return paths, rows.Err() } +// CountActiveWorktreeReferences counts non-deleted worktree associations whose +// sessions are not terminal, excluding associations owned by the caller. +func (s *SQLiteStore) CountActiveWorktreeReferences( + ctx context.Context, + worktreeID string, + excludeSessionIDs []string, +) (int, error) { + query := ` + SELECT COUNT(*) + FROM task_session_worktrees tsw + INNER JOIN task_sessions s ON s.id = tsw.session_id + WHERE tsw.worktree_id = ? + AND tsw.status <> ? + AND tsw.deleted_at IS NULL + AND s.state NOT IN (?, ?, ?) + ` + args := []interface{}{ + worktreeID, + StatusDeleted, + models.TaskSessionStateCompleted, + models.TaskSessionStateFailed, + models.TaskSessionStateCancelled, + } + if len(excludeSessionIDs) > 0 { + query += ` AND tsw.session_id NOT IN (?)` + args = append(args, excludeSessionIDs) + } + query, args, err := sqlx.In(query, args...) + if err != nil { + return 0, err + } + var count int + err = s.ro.QueryRowContext(ctx, s.ro.Rebind(query), args...).Scan(&count) + return count, err +} + // scanWorktrees is a helper to scan multiple worktree rows. func (s *SQLiteStore) scanWorktrees(rows *sql.Rows) ([]*Worktree, error) { var result []*Worktree