byre/internal/store/discard.go

package store

import (
	"context"
	"database/sql"
	"errors"
	"fmt"

	"g.fc.run/misc/byre/internal/api"
)

type DiscardState string

const (
	DiscardNone    DiscardState = ""
	DiscardHeld    DiscardState = "held"
	DiscardRunning DiscardState = "running"
	DiscardClosing DiscardState = "closing"
)

var (
	ErrDiscardPending       = errors.New("the task is already being discarded")
	ErrMerging              = errors.New("the task is merging")
	ErrMergingIntoParent    = errors.New("the task is merging into its parent task")
	ErrHasSubtasks          = errors.New("the task has active subtasks")
	ErrParentNotOpen        = errors.New("the parent task is closed, being discarded or finishing its merge")
	ErrParentMergeRequested = errors.New("the subtask is already merging into its parent")
)

type Discard struct {
	State   DiscardState
	By      string
	Session string
}

func (s *Store) GetDiscard(ctx context.Context, taskID string) (Discard, error) {
	d := Discard{}
	err := s.db.QueryRowContext(ctx, `SELECT discard_state, discarding_by, discard_session FROM tasks WHERE id = ?`, taskID).Scan(&d.State, &d.By, &d.Session)
	if errors.Is(err, sql.ErrNoRows) {
		return d, ErrNotFound
	}
	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 {
		projectID := ""
		status := ""
		mergingBy := ""
		finalizing := ""
		parentMergeBy := ""
		subtasks := 0
		err := tx.QueryRowContext(ctx, `
			SELECT project_id, status, merging_by, finalizing, parent_merge_by, discard_state,
				(SELECT count(*) FROM tasks AS sub WHERE sub.parent_task_id = tasks.id AND sub.status = ?)
			FROM tasks WHERE id = ?`, string(api.TaskStatusActive), taskID).Scan(&projectID, &status, &mergingBy, &finalizing, &parentMergeBy, &state, &subtasks)
		if errors.Is(err, sql.ErrNoRows) {
			return ErrNotFound
		}
		if err != nil {
			return err
		}
		switch {
		case status != string(api.TaskStatusActive):
			return fmt.Errorf("%w: task %s is %s", ErrNotFound, taskID, status)
		case state == DiscardClosing:
			return nil
		case state != DiscardNone:
			return ErrDiscardPending
		case mergingBy != "" || finalizing != "":
			return ErrMerging
		case parentMergeBy != "":
			return ErrMergingIntoParent
		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 {
			return err
		}
		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 {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		projectID := ""
		state := DiscardNone
		sessState := ""
		err := tx.QueryRowContext(ctx, `
			SELECT t.project_id, t.discard_state, s.state FROM tasks t JOIN agent_sessions s ON s.task_id = t.id
			WHERE t.id = ? AND s.id = ? AND s.role = ?`, taskID, sessionID, string(api.Implementer)).Scan(&projectID, &state, &sessState)
		if errors.Is(err, sql.ErrNoRows) {
			return ErrNotFound
		}
		if err != nil {
			return err
		}
		if state != DiscardHeld || sessState != string(api.AgentSessionStateIdle) {
			return fmt.Errorf("%w: discard is %q and session %s is %s", ErrStateChanged, state, sessionID, sessState)
		}
		_, err = tx.ExecContext(ctx, `UPDATE tasks SET discard_state = ?, discard_session = ? WHERE id = ?`, string(DiscardRunning), sessionID, taskID)
		if err != nil {
			return err
		}
		_, err = tx.ExecContext(ctx, `UPDATE agent_sessions SET state = ? WHERE id = ?`, string(api.AgentSessionStateRunning), sessionID)
		if err != nil {
			return err
		}
		return emit(sessionEvent("session.state", projectID, taskID, "", sessionID, actor, map[string]any{"state": string(api.AgentSessionStateRunning)}))
	})
}

func (s *Store) EndTurn(ctx context.Context, sessionID, actor string, payload map[string]any, turn *api.Event, reviewRequest, finalizeMerge bool) (DiscardState, string, error) {
	next := DiscardNone
	finalizing := ""
	err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		taskID := ""
		projectID := ""
		owner := ""
		role := ""
		state := DiscardNone
		discardSession := ""
		err := tx.QueryRowContext(ctx, `
			SELECT coalesce(t.id, ''), coalesce(t.project_id, ''), s.owner, s.role, coalesce(t.discard_state, ''), coalesce(t.discard_session, '')
			FROM agent_sessions s LEFT JOIN tasks t ON t.id = s.task_id WHERE s.id = ?`, sessionID).Scan(&taskID, &projectID, &owner, &role, &state, &discardSession)
		if errors.Is(err, sql.ErrNoRows) {
			return ErrNotFound
		}
		if err != nil {
			return err
		}

		if turn != nil {
			err = emit(*turn)
			if err != nil {
				return err
			}
		}
		if reviewRequest {
			_, err = tx.ExecContext(ctx, `UPDATE tasks SET review_requests = review_requests + 1 WHERE id = ?`, taskID)
			if err != nil {
				return err
			}
		}
		_, err = tx.ExecContext(ctx, `UPDATE agent_sessions SET state = ? WHERE id = ?`, string(api.AgentSessionStateIdle), sessionID)
		if err != nil {
			return err
		}
		if payload == nil {
			payload = map[string]any{}
		}
		payload["state"] = string(api.AgentSessionStateIdle)
		err = emit(sessionEvent("session.state", projectID, taskID, owner, sessionID, actor, payload))
		if err != nil {
			return err
		}

		if role != string(api.Implementer) {
			return nil
		}
		switch {
		case state == DiscardHeld:
			next = DiscardHeld
		case state == DiscardRunning && discardSession == sessionID:
			next = DiscardClosing
			_, err = tx.ExecContext(ctx, `UPDATE tasks SET discard_state = ? WHERE id = ?`, string(DiscardClosing), taskID)
		case state == DiscardNone && finalizeMerge:
			err = tx.QueryRowContext(ctx, `UPDATE tasks SET finalizing = merging_by, merging_by = '' WHERE id = ? AND merging_by != '' RETURNING finalizing`, taskID).Scan(&finalizing)
			if errors.Is(err, sql.ErrNoRows) {
				err = nil
			}
		}
		return err
	})
	return next, finalizing, err
}

func (s *Store) CloseDiscard(ctx context.Context, taskID string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `UPDATE tasks SET discard_state = ? WHERE id = ? AND discard_state != ''`, string(DiscardClosing), taskID)
		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))
}