Run discards on the state-driven coordinator

internal/store/coordinate.go

@@ -0,0 +1,69 @@
+package store
+
+import (
+	"context"
+	"database/sql"
+)
+
+type Work struct {
+	AbandonedMerges []string
+	StartMerges     []ClaimedMerge
+	FinalizeMerges  []FinalizingMerge
+	StartDiscards   []string
+	FinishDiscards  []string
+}
+
+func (s *Store) Coordinate(ctx context.Context) (Work, error) {
+	work := Work{}
+	err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
+		var err error
+		work.AbandonedMerges, err = abandonMerges(ctx, tx, emit)
+		if err != nil {
+			return err
+		}
+		err = claimHeldMerges(ctx, tx, emit)
+		if err != nil {
+			return err
+		}
+		work.StartMerges, err = startableMerges(ctx, tx)
+		if err != nil {
+			return err
+		}
+		work.FinalizeMerges, err = finalizingMerges(ctx, tx)
+		if err != nil {
+			return err
+		}
+		_, err = tx.ExecContext(ctx, `UPDATE tasks SET discard_state = 'closing'
+			WHERE status = 'active' AND discard_state = 'running'
+				AND NOT EXISTS (SELECT 1 FROM agent_sessions s WHERE s.id = tasks.discard_session AND s.state = 'running')`)
+		if err != nil {
+			return err
+		}
+		work.StartDiscards, err = queryIDs(ctx, tx, `SELECT id FROM tasks WHERE status = 'active' AND discard_state = 'held' AND NOT `+implementerRunning)
+		if err != nil {
+			return err
+		}
+		work.FinishDiscards, err = queryIDs(ctx, tx, `SELECT id FROM tasks WHERE status = 'active' AND discard_state = 'closing'`)
+		return err
+	})
+	return work, err
+}
+
+func queryIDs(ctx context.Context, tx *sql.Tx, query string) ([]string, error) {
+	rows, err := tx.QueryContext(ctx, query)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	ids := []string{}
+	for rows.Next() {
+		id := ""
+		err = rows.Scan(&id)
+		if err != nil {
+			return nil, err
+		}
+		ids = append(ids, id)
+	}
+	return ids, rows.Err()
+}

internal/store/discard.go

@@ -42,9 +42,9 @@ func (s *Store) GetDiscard(ctx context.Context, taskID string) (Discard, error)
 	return d, err
 }
 
-func (s *Store) RequestDiscard(ctx context.Context, taskID, actor string) (DiscardState, error) {
-	state := DiscardNone
-	err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
+func (s *Store) RequestDiscard(ctx context.Context, taskID, actor string) error {
+	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
+		state := DiscardNone
 		projectID := ""
 		status := ""
 		mergingBy := ""
@@ -75,7 +75,6 @@ func (s *Store) RequestDiscard(ctx context.Context, taskID, actor string) (Disca
 		case subtasks > 0:
 			return ErrHasSubtasks
 		}
-		state = DiscardHeld
 		_, err = tx.ExecContext(ctx, `UPDATE tasks SET discard_state = ?, discarding_by = ?, discard_session = '', merge_queued_by = '', merge_queued_at = '', merge_after_review_by = '', merge_after_review_id = '', merge_after_review_target = '' WHERE id = ?`,
 			string(DiscardHeld), actor, taskID)
 		if err != nil {
@@ -83,7 +82,6 @@ func (s *Store) RequestDiscard(ctx context.Context, taskID, actor string) (Disca
 		}
 		return emit(api.Event{Kind: "task.discard_started", ProjectId: &projectID, TaskId: &taskID, Actor: actor, Payload: map[string]any{"id": taskID}})
 	})
-	return state, err
 }
 
 func (s *Store) StartDiscardTurn(ctx context.Context, taskID, sessionID, actor string) error {
@@ -176,8 +174,7 @@ func (s *Store) EndTurn(ctx context.Context, sessionID, actor string, payload ma
 		case state == DiscardHeld:
 			end.Discard = DiscardHeld
 		case state == DiscardRunning && discardSession == sessionID:
-			end.Discard = DiscardClosing
-			_, err = tx.ExecContext(ctx, `UPDATE tasks SET discard_state = ? WHERE id = ?`, string(DiscardClosing), taskID)
+			end.Discard = DiscardRunning
 		case state == DiscardNone && finalizeMerge && lead:
 			err = tx.QueryRowContext(ctx, `UPDATE tasks SET finalizing = merging_by, merging_by = '' WHERE id = ? AND merging_by != '' AND merge_after_turn = 0 RETURNING finalizing`, taskID).Scan(&end.FinalizeBy)
 			if errors.Is(err, sql.ErrNoRows) {
@@ -195,7 +192,3 @@ func (s *Store) CloseDiscard(ctx context.Context, taskID string) error {
 		return err
 	})
 }
-
-func (s *Store) PendingDiscards(ctx context.Context) ([]api.Task, error) {
-	return s.queryTasks(ctx, `SELECT `+taskColumns+` FROM tasks WHERE status = ? AND discard_state != ''`, string(api.TaskStatusActive))
-}

internal/store/merges.go

@@ -62,34 +62,6 @@ type FinalizingMerge struct {
 	Unchanged bool
 }
 
-type MergeWork struct {
-	Abandoned []string
-	Start     []ClaimedMerge
-	Finalize  []FinalizingMerge
-}
-
-func (s *Store) CoordinateMerges(ctx context.Context) (MergeWork, error) {
-	work := MergeWork{}
-	err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
-		var err error
-		work.Abandoned, err = abandonMerges(ctx, tx, emit)
-		if err != nil {
-			return err
-		}
-		err = claimHeldMerges(ctx, tx, emit)
-		if err != nil {
-			return err
-		}
-		work.Start, err = startableMerges(ctx, tx)
-		if err != nil {
-			return err
-		}
-		work.Finalize, err = finalizingMerges(ctx, tx)
-		return err
-	})
-	return work, err
-}
-
 type idRow struct {
 	id      string
 	project string