byre/internal/store/workflow.go

package store

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

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

type FrozenWorkflow struct {
	Name         string
	DefinitionID string
	Definition   api.WorkflowDefinition
	Brief        string
	Title        string
	Audit        string
}

func (s *Store) ListWorkflows(ctx context.Context) ([]api.WorkflowSummary, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT w.name, d.description FROM workflows w JOIN workflow_definitions d ON d.id = w.definition_id ORDER BY w.name`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	workflows := []api.WorkflowSummary{}
	for rows.Next() {
		w := api.WorkflowSummary{}
		err = rows.Scan(&w.Name, &w.Description)
		if err != nil {
			return nil, err
		}
		workflows = append(workflows, w)
	}
	return workflows, rows.Err()
}

func (s *Store) GetWorkflow(ctx context.Context, name string) (api.WorkflowDefinition, error) {
	f, err := s.FreezeWorkflow(ctx, name)
	return f.Definition, err
}

var ErrWorkflowExists = errors.New("workflow exists")

func (s *Store) PutWorkflow(ctx context.Context, name, from string, def api.WorkflowDefinition, actor string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		if from != "" {
			found := false
			err := tx.QueryRowContext(ctx, `SELECT EXISTS (SELECT 1 FROM workflows WHERE name = ?)`, from).Scan(&found)
			if err != nil {
				return err
			}
			if !found {
				return fmt.Errorf("%w: no workflow %q", ErrNotFound, from)
			}
		}
		if from != "" && from != name {
			taken := false
			err := tx.QueryRowContext(ctx, `SELECT EXISTS (SELECT 1 FROM workflows WHERE name = ?)`, name).Scan(&taken)
			if err != nil {
				return err
			}
			if taken {
				return fmt.Errorf("%w: a workflow named %q already exists", ErrWorkflowExists, name)
			}
			_, err = tx.ExecContext(ctx, `DELETE FROM workflows WHERE name = ?`, from)
			if err != nil {
				return err
			}
			err = repointProjects(ctx, tx, emit, from, name, actor)
			if err != nil {
				return err
			}
			err = emit(api.Event{Kind: "workflow.deleted", Actor: actor, Payload: map[string]any{"name": from}})
			if err != nil {
				return err
			}
		}
		defID, err := insertDefinition(ctx, tx, def)
		if err != nil {
			return err
		}
		_, err = tx.ExecContext(ctx, `INSERT INTO workflows (name, definition_id) VALUES (?, ?)
			ON CONFLICT (name) DO UPDATE SET definition_id = excluded.definition_id`, name, defID)
		if err != nil {
			return err
		}
		return emit(api.Event{Kind: "workflow.updated", Actor: actor, Payload: map[string]any{"name": name}})
	})
}

func repointProjects(ctx context.Context, tx *sql.Tx, emit emitFunc, from, to, actor string) error {
	rows, err := tx.QueryContext(ctx, `
		UPDATE projects SET
			review_workflow = CASE WHEN review_workflow = ?1 THEN ?2 ELSE review_workflow END,
			audit_workflow = CASE WHEN audit_workflow = ?1 THEN ?2 ELSE audit_workflow END
		WHERE review_workflow = ?1 OR audit_workflow = ?1
		RETURNING `+projectColumns, from, to)
	if err != nil {
		return err
	}
	projects := []api.Project{}
	for rows.Next() {
		p, err := scanProject(rows)
		if err != nil {
			rows.Close()
			return err
		}
		projects = append(projects, p)
	}
	err = errors.Join(rows.Err(), rows.Close())
	if err != nil {
		return err
	}
	for _, p := range projects {
		err = emit(api.Event{Kind: "project.updated", ProjectId: &p.Id, Actor: actor, Payload: ToPayload(p)})
		if err != nil {
			return err
		}
	}
	return nil
}

func (s *Store) DeleteWorkflow(ctx context.Context, name, actor string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		res, err := tx.ExecContext(ctx, `DELETE FROM workflows WHERE name = ?`, name)
		if err != nil {
			return err
		}
		n, err := res.RowsAffected()
		if err != nil {
			return err
		}
		if n == 0 {
			return fmt.Errorf("%w: no workflow %q", ErrNotFound, name)
		}
		err = repointProjects(ctx, tx, emit, name, "", actor)
		if err != nil {
			return err
		}
		return emit(api.Event{Kind: "workflow.deleted", Actor: actor, Payload: map[string]any{"name": name}})
	})
}

func (s *Store) FreezeWorkflow(ctx context.Context, name string) (FrozenWorkflow, error) {
	f := FrozenWorkflow{Name: name}
	err := s.db.QueryRowContext(ctx, `SELECT definition_id FROM workflows WHERE name = ?`, name).Scan(&f.DefinitionID)
	if errors.Is(err, sql.ErrNoRows) {
		return FrozenWorkflow{}, fmt.Errorf("%w: no workflow %q", ErrNotFound, name)
	}
	if err != nil {
		return FrozenWorkflow{}, err
	}
	f.Definition, err = loadDefinition(ctx, s.db, f.DefinitionID)
	return f, err
}

func (s *Store) TaskWorkflow(ctx context.Context, taskID string) (FrozenWorkflow, error) {
	f := FrozenWorkflow{}
	defID := sql.NullString{}
	err := s.db.QueryRowContext(ctx, `
		SELECT r.workflow_name, r.workflow_definition_id, r.workflow_brief, r.workflow_title, r.audit
		FROM tasks t JOIN tasks r ON r.id = CASE WHEN t.workflow_stage != '' THEN t.parent_task_id ELSE t.id END
		WHERE t.id = ?`, taskID).Scan(&f.Name, &defID, &f.Brief, &f.Title, &f.Audit)
	if errors.Is(err, sql.ErrNoRows) {
		return FrozenWorkflow{}, fmt.Errorf("%w: no task %s", ErrNotFound, taskID)
	}
	if err != nil {
		return FrozenWorkflow{}, err
	}
	if !defID.Valid {
		return FrozenWorkflow{}, fmt.Errorf("%w: task %s has no workflow", ErrNotFound, taskID)
	}
	f.DefinitionID = defID.String
	f.Definition, err = loadDefinition(ctx, s.db, f.DefinitionID)
	return f, err
}

func (s *Store) LastAudit(ctx context.Context, projectID string) (string, bool, error) {
	commit := ""
	active := false
	err := s.db.QueryRowContext(ctx, `
		SELECT
			coalesce((SELECT audit FROM tasks WHERE project_id = ?1 AND audit != '' ORDER BY rowid DESC LIMIT 1), ''),
			EXISTS (SELECT 1 FROM tasks WHERE project_id = ?1 AND audit != '' AND status = ?2)
	`, projectID, string(api.TaskStatusActive)).Scan(&commit, &active)
	return commit, active, err
}

func (s *Store) RunTasks(ctx context.Context, run string) ([]api.Task, error) {
	return s.queryTasks(ctx, `SELECT `+taskColumns+` FROM tasks WHERE parent_task_id = ? AND workflow_stage != '' ORDER BY created_at`, run)
}

func (s *Store) RunEvents(ctx context.Context, run string) ([]api.Event, error) {
	rows, err := s.db.QueryContext(ctx, `
		SELECT `+eventColumns+` FROM events
		WHERE task_id IN (SELECT id FROM tasks WHERE parent_task_id = ?1 AND workflow_stage != '') AND kind IN ('task.workflow_stage', 'task.discarded')
		ORDER BY seq
	`, run)
	if err != nil {
		return nil, err
	}
	return scanEvents(rows)
}

func (s *Store) RecordStage(ctx context.Context, taskID, output, commit string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `UPDATE tasks SET stage_output = ?, stage_commit = ?, ready_prompt_seq = 0 WHERE id = ?`, output, commit, taskID)
		if err != nil {
			return err
		}
		return settleTaskState(ctx, tx, emit, taskID, "byre")
	})
}

func (s *Store) StageResult(ctx context.Context, taskID string) (string, string, error) {
	output := ""
	commit := ""
	err := s.db.QueryRowContext(ctx, `SELECT stage_output, stage_commit FROM tasks WHERE id = ?`, taskID).Scan(&output, &commit)
	return output, commit, err
}

func (s *Store) SetStageOutput(ctx context.Context, sessionID, output string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `UPDATE agent_sessions SET stage_output = ? WHERE id = ?`, output, sessionID)
		return err
	})
}

func (s *Store) FinishStage(ctx context.Context, sessionID string) (bool, error) {
	finished := false
	err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		res, err := tx.ExecContext(ctx, `UPDATE agent_sessions SET stage_finished = 1 WHERE id = ? AND stage_output != '' AND NOT EXISTS (SELECT 1 FROM tasks WHERE tasks.id = agent_sessions.task_id AND tasks.discard_state != '')`, sessionID)
		if err != nil {
			return err
		}
		n, err := res.RowsAffected()
		finished = n > 0
		return err
	})
	return finished, err
}

func (s *Store) DropStageOutput(ctx context.Context, sessionID string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `UPDATE agent_sessions SET stage_output = '' WHERE id = ? AND stage_finished = 0`, sessionID)
		return err
	})
}

func (s *Store) StageOutput(ctx context.Context, sessionID string) (string, error) {
	output := ""
	err := s.db.QueryRowContext(ctx, `SELECT stage_output FROM agent_sessions WHERE id = ?`, sessionID).Scan(&output)
	return output, err
}

func (s *Store) FinishedStages(ctx context.Context) ([]string, error) {
	return queryIDs(ctx, s.db, `SELECT id FROM agent_sessions WHERE stage_finished = 1 ORDER BY created_at`)
}

func (s *Store) ClearStage(ctx context.Context, sessionID string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		return clearStage(ctx, tx, sessionID)
	})
}

func clearStage(ctx context.Context, tx *sql.Tx, sessionID string) error {
	_, err := tx.ExecContext(ctx, `UPDATE agent_sessions SET stage_output = '', stage_finished = 0 WHERE id = ?`, sessionID)
	return err
}

func (s *Store) FinishWorkflow(ctx context.Context, run api.Task, sessionID, workflowName, output string, kept []string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		err := clearStage(ctx, tx, sessionID)
		if err != nil {
			return err
		}
		return emit(api.Event{
			Kind:      "task.workflow_finished",
			ProjectId: &run.ProjectId,
			TaskId:    &run.Id,
			Actor:     "byre",
			Payload:   map[string]any{"id": run.Id, "workflow": workflowName, "output": output, "kept": kept},
		})
	})
}

func (s *Store) WorkflowFinished(ctx context.Context, run string) (bool, error) {
	finished := false
	err := s.db.QueryRowContext(ctx, `SELECT EXISTS (SELECT 1 FROM events WHERE task_id = ? AND kind = 'task.workflow_finished')`, run).Scan(&finished)
	return finished, err
}