byre/internal/store/queue.go

package store

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"time"

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

type QueuedTurn struct {
	SessionID string
	TaskID    string
	Vendor    api.AgentSessionVendor
	Actor     string
	Text      string
	Images    []string
	CreatedAt time.Time
}

func (s *Store) SessionStarted(ctx context.Context, sessionID string) (bool, error) {
	started := false
	err := s.db.QueryRowContext(ctx, `SELECT EXISTS (SELECT 1 FROM events WHERE agent_session_id = ? AND kind = 'user.message')`, sessionID).Scan(&started)
	return started, err
}

func (s *Store) CountRunning(ctx context.Context) (map[api.AgentSessionVendor]int, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT vendor, COUNT(*) FROM agent_sessions WHERE state = ? AND owner = '' GROUP BY vendor`, string(api.AgentSessionStateRunning))
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	counts := map[api.AgentSessionVendor]int{}
	for rows.Next() {
		vendor := ""
		n := 0
		err = rows.Scan(&vendor, &n)
		if err != nil {
			return nil, err
		}
		counts[api.AgentSessionVendor(vendor)] = n
	}
	return counts, rows.Err()
}

func (s *Store) EnqueueTurn(ctx context.Context, q QueuedTurn) error {
	images, err := json.Marshal(append([]string{}, q.Images...))
	if err != nil {
		return err
	}
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `INSERT INTO turn_queue (session_id, actor, text, images, created_at) VALUES (?, ?, ?, ?, ?)`,
			q.SessionID, q.Actor, q.Text, string(images), formatTime(q.CreatedAt))
		return err
	})
}

func (s *Store) QueuedTurns(ctx context.Context) ([]QueuedTurn, error) {
	rows, err := s.db.QueryContext(ctx, `
		SELECT q.session_id, s.task_id, s.vendor, q.actor, q.text, q.images, q.created_at
		FROM turn_queue q JOIN agent_sessions s ON s.id = q.session_id
		ORDER BY q.created_at
	`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	turns := []QueuedTurn{}
	for rows.Next() {
		q := QueuedTurn{}
		vendor := ""
		images := ""
		createdAt := ""
		err = rows.Scan(&q.SessionID, &q.TaskID, &vendor, &q.Actor, &q.Text, &images, &createdAt)
		if err != nil {
			return nil, err
		}
		err = json.Unmarshal([]byte(images), &q.Images)
		if err != nil {
			return nil, err
		}
		q.Vendor = api.AgentSessionVendor(vendor)
		q.CreatedAt, err = parseTime(createdAt)
		if err != nil {
			return nil, err
		}
		turns = append(turns, q)
	}
	return turns, rows.Err()
}

func (s *Store) DeleteQueuedTurn(ctx context.Context, sessionID string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		res, err := tx.ExecContext(ctx, `DELETE FROM turn_queue WHERE session_id = ?`, sessionID)
		if err != nil {
			return err
		}
		n, err := res.RowsAffected()
		if err != nil {
			return err
		}
		if n == 0 {
			return errors.Join(ErrNotFound, errors.New("no queued turn for session "+sessionID))
		}
		return nil
	})
}