byre/internal/store/readyprompt.go

package store

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

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

type ReadyPromptKind string

const (
	ReadyPromptReady       ReadyPromptKind = "ready"
	ReadyPromptUnchanged   ReadyPromptKind = "unchanged"
	ReadyPromptFinishStage ReadyPromptKind = "finish_stage"
)

type ReadyPrompt struct {
	TaskID    string
	SessionID string
	Kind      ReadyPromptKind
	Snapshot  int
}

func (s *Store) SetReadyPromptSeq(ctx context.Context, id string, seq int64) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `UPDATE tasks SET ready_prompt_seq = ? WHERE id = ?`, seq, id)
		return err
	})
}

func readyPrompts(ctx context.Context, tx *sql.Tx, emit emitFunc) ([]ReadyPrompt, error) {
	type candidate struct {
		ReadyPrompt
		projectID  string
		stage      string
		run        string
		output     string
		background bool
		workflow   string
	}
	rows, err := tx.QueryContext(ctx, `
		SELECT tasks.id, tasks.project_id, s.id, tasks.workflow_stage, coalesce(tasks.parent_task_id, ''), tasks.stage_output, `+backgroundExpr+`,
			(SELECT p.review_workflow FROM projects p WHERE p.id = tasks.project_id)
		FROM tasks JOIN agent_sessions s ON s.id = (
			SELECT l.id FROM agent_sessions l
			WHERE l.task_id = tasks.id AND l.role != 'reviewer' AND l.state != 'completed'
				AND l.member = (SELECT f.member FROM agent_sessions f WHERE f.task_id = tasks.id AND f.role != 'reviewer' ORDER BY f.created_at LIMIT 1)
			ORDER BY l.created_at DESC LIMIT 1)
		WHERE tasks.status = 'active' AND tasks.state IN ('idle', 'finished', 'ready_to_merge', 'discard_suggested') AND s.state = 'idle'
			AND NOT EXISTS (SELECT 1 FROM inbox_items i WHERE i.task_id = tasks.id AND i.state = 'open')
			AND NOT EXISTS (SELECT 1 FROM task_inputs q WHERE q.task_id = tasks.id)
			AND EXISTS (SELECT 1 FROM events e WHERE e.task_id = tasks.id AND e.seq > tasks.ready_prompt_seq
				AND ((e.kind = 'snapshot.created' AND json_extract(e.payload, '$.trigger') NOT IN `+cutPointTriggers+`)
					OR (e.kind = 'user.message' AND e.actor != ? AND e.actor NOT LIKE '% (lead)' AND json_extract(e.payload, '$.bounce') IS NOT 1)))`, agentActor)
	if err != nil {
		return nil, err
	}
	found := []candidate{}
	for rows.Next() {
		c := candidate{}
		err = rows.Scan(&c.TaskID, &c.projectID, &c.SessionID, &c.stage, &c.run, &c.output, &c.background, &c.workflow)
		if err != nil {
			rows.Close()
			return nil, err
		}
		found = append(found, c)
	}
	err = errors.Join(rows.Err(), rows.Close())
	if err != nil {
		return nil, err
	}

	prompts := []ReadyPrompt{}
	for _, c := range found {
		stopped, err := sessionStopped(ctx, tx, c.TaskID, c.SessionID)
		if err != nil {
			return nil, err
		}
		if stopped {
			continue
		}
		c.Kind = ReadyPromptReady
		if c.stage != "" {
			frozen, err := taskWorkflow(ctx, tx, c.TaskID)
			if err != nil {
				return nil, err
			}
			if c.output != "" && len(workflow.Next(frozen.Definition, c.stage)) > 0 {
				continue
			}
			finished := false
			err = tx.QueryRowContext(ctx, `SELECT EXISTS (SELECT 1 FROM events WHERE task_id = ? AND kind = 'task.workflow_finished')`, c.run).Scan(&finished)
			if err != nil {
				return nil, err
			}
			if c.output == "" && !finished {
				c.Kind = ReadyPromptFinishStage
			}
		}

		settle := ""
		err = tx.QueryRowContext(ctx, `SELECT n, settle_state FROM snapshots WHERE task_id = ? AND trigger NOT IN `+cutPointTriggers+` ORDER BY n DESC LIMIT 1`, c.TaskID).Scan(&c.Snapshot, &settle)
		if errors.Is(err, sql.ErrNoRows) {
			continue
		}
		if err != nil {
			return nil, err
		}
		reviews, err := snapshotReviews(ctx, tx, c.TaskID, c.Snapshot)
		if err != nil {
			return nil, err
		}
		if len(reviews) == 0 && c.workflow != "" || len(reviews) > 0 && settle != "done" {
			continue
		}
		blocked := false
		for _, r := range reviews {
			if r.State == api.ReviewStateFailed && c.background {
				err = emit(api.Event{Kind: "session.error", ProjectId: &c.projectID, TaskId: &c.TaskID, Actor: "byre", Payload: map[string]any{
					"error": fmt.Sprintf("Review %q of snapshot %d failed: %s. The audit can't go on without you.", r.Stage, c.Snapshot, r.Error),
				}})
				if err != nil {
					return nil, err
				}
			}
			if r.State != api.ReviewStateDone && r.State != api.ReviewStateSkipped {
				blocked = true
				break
			}
		}
		if blocked {
			continue
		}
		if c.Kind == ReadyPromptReady && len(reviews) == 1 && reviews[0].State == api.ReviewStateSkipped {
			c.Kind = ReadyPromptUnchanged
		}
		prompts = append(prompts, c.ReadyPrompt)
	}
	return prompts, nil
}