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 := "" lead := false err := tx.QueryRowContext(ctx, ` SELECT coalesce(t.id, ''), coalesce(t.project_id, ''), s.owner, s.role, coalesce(t.discard_state, ''), coalesce(t.discard_session, ''), s.member = coalesce((SELECT f.member FROM agent_sessions f WHERE f.task_id = s.task_id AND f.role = s.role ORDER BY f.created_at LIMIT 1), '') 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, &lead) 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 && lead: 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)) }