112 lines
2.8 KiB
Go
112 lines
2.8 KiB
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 = ? 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
|
|
})
|
|
}
|