From cbf35e875cc6a0b8d4fb4c8d541c1fc8e60c0d89 Mon Sep 17 00:00:00 2001 From: Aditi Chauhan Date: Tue, 30 Jun 2026 14:38:02 +0530 Subject: [PATCH] fix: preserve dirty workspace project kills --- backend/internal/domain/project.go | 8 +-- backend/internal/session_manager/manager.go | 16 +++-- .../internal/session_manager/manager_test.go | 58 ++++++++++++------- .../sqlite/store/session_worktree_store.go | 13 ++--- 4 files changed, 57 insertions(+), 38 deletions(-) diff --git a/backend/internal/domain/project.go b/backend/internal/domain/project.go index fd0fc25e2..f99162ae5 100644 --- a/backend/internal/domain/project.go +++ b/backend/internal/domain/project.go @@ -56,10 +56,8 @@ type SessionWorktreeRecord struct { BaseSHA string WorktreePath string PreservedRef string - // ponytail: State mirrors session_worktrees.state, an enum that is unused - // multi-repo scaffolding. The save/restore lifecycle reads and writes only - // PreservedRef and row presence; State is never set by any live code path - // and always resolves to the column default ('active' on insert). Wire it - // when multi-repo worktree lifecycle states actually ship. + // State mirrors session_worktrees.state. Workspace project lifecycle code + // uses it to distinguish live rows from shutdown-saved removed rows and + // rows that could not be removed cleanly. State string } diff --git a/backend/internal/session_manager/manager.go b/backend/internal/session_manager/manager.go index 45d583818..897f01b0e 100644 --- a/backend/internal/session_manager/manager.go +++ b/backend/internal/session_manager/manager.go @@ -517,7 +517,7 @@ func (m *Manager) Kill(ctx context.Context, id domain.SessionID) (bool, error) { if rows, ok, rowErr := m.workspaceProjectRows(ctx, rec); rowErr != nil { return false, fmt.Errorf("kill %s: workspace rows: %w", id, rowErr) } else if ok { - cleaned, err := m.destroyWorkspaceProjectRows(ctx, rows) + cleaned, err := m.destroyWorkspaceProjectRows(ctx, rows, false) if err != nil { if errors.Is(err, ports.ErrWorkspaceDirty) { return false, nil @@ -986,7 +986,12 @@ func (m *Manager) sessionWorktreeRowsToRepoInfos(ctx context.Context, project do for _, row := range rows { repoPath := repoPaths[row.RepoName] if repoPath == "" { - return nil, fmt.Errorf("repo %q is not registered for workspace project %s", row.RepoName, project.ID) + m.logger.WarnContext(ctx, "session worktree row references unregistered workspace repo", + "sessionID", rec.ID, + "projectID", project.ID, + "repo", row.RepoName, + ) + continue } out = append(out, ports.WorkspaceRepoInfo{ RepoName: row.RepoName, @@ -1041,7 +1046,7 @@ func (m *Manager) saveAndTeardownWorkspaceProject(ctx context.Context, rec domai return nil } -func (m *Manager) destroyWorkspaceProjectRows(ctx context.Context, rows []ports.WorkspaceRepoInfo) (bool, error) { +func (m *Manager) destroyWorkspaceProjectRows(ctx context.Context, rows []ports.WorkspaceRepoInfo, preserveDirty bool) (bool, error) { adapter, ok := m.workspace.(ports.WorkspaceProject) if !ok { return false, errors.New("workspace project lifecycle is not supported by workspace adapter") @@ -1055,6 +1060,9 @@ func (m *Manager) destroyWorkspaceProjectRows(ctx context.Context, rows []ports. if err := adapter.DestroyWorkspaceProjectWorktree(ctx, rows[i]); err != nil { preservedRef := "" if errors.Is(err, ports.ErrWorkspaceDirty) { + if !preserveDirty { + return cleaned, err + } ref, preserveErr := adapter.StashWorkspaceProjectWorktree(ctx, rows[i]) if preserveErr == nil { preservedRef = ref @@ -1200,7 +1208,7 @@ func (m *Manager) Cleanup(ctx context.Context, project domain.ProjectID) (Cleanu result.Skipped = append(result.Skipped, CleanupSkip{SessionID: rec.ID, Reason: "workspace teardown failed"}) continue } else if ok { - if _, err := m.destroyWorkspaceProjectRows(ctx, rows); err != nil { + if _, err := m.destroyWorkspaceProjectRows(ctx, rows, true); err != nil { if !errors.Is(err, ports.ErrWorkspaceDirty) { m.logger.Warn("cleanup: workspace teardown failed", "sessionID", rec.ID, "path", ws.Path, "error", err) } diff --git a/backend/internal/session_manager/manager_test.go b/backend/internal/session_manager/manager_test.go index d40907a65..d9b472625 100644 --- a/backend/internal/session_manager/manager_test.go +++ b/backend/internal/session_manager/manager_test.go @@ -814,7 +814,33 @@ func TestKill_WorkspaceProjectDestroysChildrenBeforeRoot(t *testing.T) { } } -func TestKill_WorkspaceProjectDirtyRowsArePreservedAndForceRemoved(t *testing.T) { +func TestKill_WorkspaceProjectSkipsUnregisteredChildRows(t *testing.T) { + m, st, _, ws := newManager() + st.projects["mer"] = domain.ProjectRecord{ID: "mer", Path: "/repo/mer", Kind: domain.ProjectKindWorkspace, Config: testRoleAgents()} + st.workspaceRepo["mer"] = []domain.WorkspaceRepoRecord{{Name: "api", RelativePath: "api"}} + st.sessions["mer-1"] = domain.SessionRecord{ + ID: "mer-1", + ProjectID: "mer", + Metadata: domain.SessionMetadata{WorkspacePath: "/ws/mer-1", Branch: "ao/mer-1", RuntimeHandleID: "h1"}, + Activity: domain.Activity{State: domain.ActivityActive}, + } + st.worktrees["mer-1"] = []domain.SessionWorktreeRecord{ + {SessionID: "mer-1", RepoName: domain.RootWorkspaceRepoName, Branch: "ao/mer-1", WorktreePath: "/ws/mer-1"}, + {SessionID: "mer-1", RepoName: "old-api", Branch: "ao/mer-1", WorktreePath: "/ws/mer-1/old-api"}, + {SessionID: "mer-1", RepoName: "api", Branch: "ao/mer-1", WorktreePath: "/ws/mer-1/api"}, + } + + freed, err := m.Kill(ctx, "mer-1") + if err != nil || !freed { + t.Fatalf("freed=%v err=%v", freed, err) + } + want := []string{"DestroyWorkspaceProjectWorktree:api", "DestroyWorkspaceProjectWorktree:__root__"} + if got := ws.calls; strings.Join(got, ",") != strings.Join(want, ",") { + t.Fatalf("destroy calls = %v, want %v", got, want) + } +} + +func TestKill_WorkspaceProjectDirtyRowRefusesRemoval(t *testing.T) { m, st, _, ws := newManager() ws.destroyErr = fmt.Errorf("dirty: %w", ports.ErrWorkspaceDirty) ws.stashRef = "refs/ao/preserved/mer-1" @@ -832,31 +858,23 @@ func TestKill_WorkspaceProjectDirtyRowsArePreservedAndForceRemoved(t *testing.T) } freed, err := m.Kill(ctx, "mer-1") - if err != nil || !freed { - t.Fatalf("freed=%v err=%v, want dirty rows preserved and removed", freed, err) + if err != nil { + t.Fatalf("kill dirty workspace project err = %v, want nil", err) } - want := []string{ - "DestroyWorkspaceProjectWorktree:api", - "StashWorkspaceProjectWorktree:api", - "ForceDestroyWorkspaceProjectWorktree:api", - "DestroyWorkspaceProjectWorktree:__root__", - "StashWorkspaceProjectWorktree:__root__", - "ForceDestroyWorkspaceProjectWorktree:__root__", + if freed { + t.Fatal("freed = true, want false for preserved workspace project") } + want := []string{"DestroyWorkspaceProjectWorktree:api"} if got := ws.calls; strings.Join(got, ",") != strings.Join(want, ",") { t.Fatalf("calls = %v, want %v", got, want) } - refs := map[string]string{} - states := map[string]string{} + if !st.sessions["mer-1"].IsTerminated { + t.Fatal("session should be terminated") + } for _, row := range st.worktrees["mer-1"] { - refs[row.RepoName] = row.PreservedRef - states[row.RepoName] = row.State - } - if refs["api"] != "refs/ao/preserved/mer-1/api" || refs[domain.RootWorkspaceRepoName] != "refs/ao/preserved/mer-1/__root__" { - t.Fatalf("preserved refs = %v", refs) - } - if states["api"] != "unavailable" || states[domain.RootWorkspaceRepoName] != "unavailable" { - t.Fatalf("states = %v, want unavailable rows so RestoreAll does not resurrect killed session", states) + if row.PreservedRef != "" || row.State != "" { + t.Fatalf("dirty kill should not mutate worktree row %+v", row) + } } } diff --git a/backend/internal/storage/sqlite/store/session_worktree_store.go b/backend/internal/storage/sqlite/store/session_worktree_store.go index 9b6e1a141..b57043791 100644 --- a/backend/internal/storage/sqlite/store/session_worktree_store.go +++ b/backend/internal/storage/sqlite/store/session_worktree_store.go @@ -14,12 +14,9 @@ import ( func (s *Store) UpsertSessionWorktree(ctx context.Context, row domain.SessionWorktreeRecord) error { s.writeMu.Lock() defer s.writeMu.Unlock() - // ponytail: session_worktrees.state is unused multi-repo scaffolding; no - // live code path sets domain.SessionWorktreeRecord.State, so it arrives - // here as "". The generated upsert includes state in the INSERT column list - // and the CHECK constraint rejects "". Default to 'active' (the column - // default) so the row stays valid without touching the schema or gen code. - // Wire a real value when multi-repo worktree lifecycle states ship. + // An empty state means the caller is creating or updating a live worktree + // row. Lifecycle paths set explicit states such as removed, unavailable, or + // retry_remove. state := row.State if state == "" { state = "active" @@ -75,8 +72,6 @@ func sessionWorktreeFromGen(row gen.SessionWorktree) domain.SessionWorktreeRecor BaseSHA: row.BaseSha, WorktreePath: row.WorktreePath, PreservedRef: row.PreservedRef, - // ponytail: state is read back from the DB but no caller uses it; - // it is unused multi-repo scaffolding (see UpsertSessionWorktree above). - State: row.State, + State: row.State, } }