474 lines
12 KiB
Go
474 lines
12 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"path/filepath"
|
|
"strings"
|
|
"time"
|
|
|
|
"g.fc.run/misc/byre/internal/api"
|
|
)
|
|
|
|
func DeliverableResponse(t api.InboxResponseType) bool {
|
|
switch t {
|
|
case api.Option, api.Other, api.Run:
|
|
return true
|
|
}
|
|
return false
|
|
}
|
|
|
|
func (s *Store) CreateInboxItem(ctx context.Context, item api.InboxItem) error {
|
|
return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
|
|
payload, err := json.Marshal(item.Payload)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
_, err = tx.ExecContext(ctx, `
|
|
INSERT INTO inbox_items (id, project_id, task_id, agent_session_id, thread_id, kind, payload, state, created_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
`, item.Id, item.ProjectId, item.TaskId, item.AgentSessionId, item.ThreadId, string(item.Kind),
|
|
string(payload), string(item.State), formatTime(item.CreatedAt))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return emitItem(ctx, tx, emit, "inbox.created", item.Id, agentActor)
|
|
})
|
|
}
|
|
|
|
const agentActor = "agent"
|
|
|
|
func (s *Store) GetInboxItem(ctx context.Context, id string) (api.InboxItem, error) {
|
|
item, err := getItem(ctx, s.db, id)
|
|
if errors.Is(err, sql.ErrNoRows) {
|
|
return item, ErrNotFound
|
|
}
|
|
return item, err
|
|
}
|
|
|
|
type InboxFilter struct {
|
|
ProjectID string
|
|
TaskID string
|
|
SessionID string
|
|
State api.InboxState
|
|
}
|
|
|
|
func (s *Store) ListInbox(ctx context.Context, f InboxFilter) ([]api.InboxItem, error) {
|
|
where := []string{"state = ?"}
|
|
args := []any{string(f.State)}
|
|
|
|
if f.ProjectID != "" {
|
|
where = append(where, "project_id = ?")
|
|
args = append(args, f.ProjectID)
|
|
}
|
|
if f.TaskID != "" {
|
|
where = append(where, "task_id = ?")
|
|
args = append(args, f.TaskID)
|
|
}
|
|
if f.SessionID != "" {
|
|
where = append(where, "agent_session_id = ?")
|
|
args = append(args, f.SessionID)
|
|
}
|
|
|
|
rows, err := s.db.QueryContext(ctx, `SELECT id FROM inbox_items WHERE `+strings.Join(where, " AND ")+` ORDER BY created_at`, args...)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
ids := []string{}
|
|
for rows.Next() {
|
|
id := ""
|
|
err = rows.Scan(&id)
|
|
if err != nil {
|
|
rows.Close()
|
|
return nil, err
|
|
}
|
|
ids = append(ids, id)
|
|
}
|
|
rows.Close()
|
|
err = rows.Err()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
items := []api.InboxItem{}
|
|
for _, id := range ids {
|
|
item, err := getItem(ctx, s.db, id)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, item)
|
|
}
|
|
return items, nil
|
|
}
|
|
|
|
type Resolution struct {
|
|
Response api.InboxResponse
|
|
State api.InboxState
|
|
}
|
|
|
|
func (s *Store) Respond(ctx context.Context, actor string, resolutions []Resolution) error {
|
|
return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
|
|
for _, r := range resolutions {
|
|
resp := r.Response
|
|
res, err := tx.ExecContext(ctx, `UPDATE inbox_items SET state = ?, resolved_by = ?, resolved_at = ? WHERE id = ? AND state = ?`,
|
|
string(r.State), actor, now(), resp.ItemId, string(api.Open))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
n, err := res.RowsAffected()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if n == 0 {
|
|
return fmt.Errorf("%w: %s", ErrNotOpen, resp.ItemId)
|
|
}
|
|
|
|
attachments, err := json.Marshal(resp.Attachments)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
images, err := json.Marshal(resp.Images)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
var options *string
|
|
if resp.Options != nil {
|
|
b, err := json.Marshal(*resp.Options)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
s := string(b)
|
|
options = &s
|
|
}
|
|
|
|
_, err = tx.ExecContext(ctx, `
|
|
INSERT INTO inbox_responses (id, item_id, user_login, type, options, text, images, attachments, created_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
`, resp.Id, resp.ItemId, resp.User, string(resp.Type), options, resp.Text, string(images), string(attachments), formatTime(resp.CreatedAt))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
err = emitItem(ctx, tx, emit, "inbox.updated", resp.ItemId, actor)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
})
|
|
}
|
|
|
|
var ErrNotOpen = errors.New("inbox item is not open")
|
|
|
|
func (s *Store) AddRunResponse(ctx context.Context, resp api.InboxResponse) (bool, error) {
|
|
added := false
|
|
err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
|
|
res, err := tx.ExecContext(ctx, `
|
|
INSERT INTO inbox_responses (id, item_id, user_login, type, text, client, exit_code, images, attachments, created_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, '[]', '[]', ?)
|
|
ON CONFLICT (id) DO NOTHING
|
|
`, resp.Id, resp.ItemId, resp.User, string(api.Run), resp.Text, *resp.Client, *resp.ExitCode, formatTime(resp.CreatedAt))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
n, err := res.RowsAffected()
|
|
if err != nil || n == 0 {
|
|
return err
|
|
}
|
|
added = true
|
|
return emitItem(ctx, tx, emit, "inbox.updated", resp.ItemId, resp.User)
|
|
})
|
|
return added, err
|
|
}
|
|
|
|
func (s *Store) BundlePath(itemID string) string {
|
|
return filepath.Join(s.dataDir(), "bundles", itemID+".tar.gz")
|
|
}
|
|
|
|
func (s *Store) WithdrawItem(ctx context.Context, id string) error {
|
|
return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
|
|
res, err := tx.ExecContext(ctx, `UPDATE inbox_items SET state = ?, resolved_at = ? WHERE id = ? AND state = ?`,
|
|
string(api.Dismissed), now(), id, string(api.Open))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
n, err := res.RowsAffected()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if n == 0 {
|
|
return ErrNotOpen
|
|
}
|
|
return emitItem(ctx, tx, emit, "inbox.updated", id, agentActor)
|
|
})
|
|
}
|
|
|
|
func dismissTaskInbox(ctx context.Context, tx *sql.Tx, emit emitFunc, taskID, actor string) error {
|
|
ids := []string{}
|
|
rows, err := tx.QueryContext(ctx, `SELECT id FROM inbox_items WHERE task_id = ? AND state = ?`, taskID, string(api.Open))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for rows.Next() {
|
|
id := ""
|
|
err = rows.Scan(&id)
|
|
if err != nil {
|
|
rows.Close()
|
|
return err
|
|
}
|
|
ids = append(ids, id)
|
|
}
|
|
rows.Close()
|
|
|
|
for _, id := range ids {
|
|
_, err = tx.ExecContext(ctx, `UPDATE inbox_items SET state = ?, resolved_by = ?, resolved_at = ? WHERE id = ?`,
|
|
string(api.Dismissed), actor, now(), id)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
err = emitItem(ctx, tx, emit, "inbox.updated", id, actor)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *Store) AwaitingHuman(ctx context.Context, sessionID string) (int, error) {
|
|
n := 0
|
|
err := s.db.QueryRowContext(ctx, `
|
|
SELECT COUNT(*) FROM inbox_items i
|
|
WHERE i.agent_session_id = ? AND i.kind = ? AND i.state = ?
|
|
AND NOT EXISTS (SELECT 1 FROM inbox_responses r WHERE r.item_id = i.id AND r.type IN (?, ?))
|
|
`, sessionID, string(api.Question), string(api.Open),
|
|
string(api.Option), string(api.Other)).Scan(&n)
|
|
return n, err
|
|
}
|
|
|
|
type Pending struct {
|
|
Item api.InboxItem
|
|
Response api.InboxResponse
|
|
}
|
|
|
|
func (s *Store) Undelivered(ctx context.Context, sessionID string) ([]Pending, error) {
|
|
rows, err := s.db.QueryContext(ctx, `
|
|
SELECT r.id, r.item_id FROM inbox_responses r JOIN inbox_items i ON i.id = r.item_id
|
|
WHERE i.agent_session_id = ? AND r.delivered_at IS NULL AND r.type IN (?, ?, ?)
|
|
ORDER BY r.created_at
|
|
`, sessionID, string(api.Option), string(api.Other), string(api.Run))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
type pair struct{ resp, item string }
|
|
pairs := []pair{}
|
|
for rows.Next() {
|
|
p := pair{}
|
|
err = rows.Scan(&p.resp, &p.item)
|
|
if err != nil {
|
|
rows.Close()
|
|
return nil, err
|
|
}
|
|
pairs = append(pairs, p)
|
|
}
|
|
rows.Close()
|
|
err = rows.Err()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
pending := []Pending{}
|
|
for _, p := range pairs {
|
|
item, err := getItem(ctx, s.db, p.item)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
for _, r := range item.Responses {
|
|
if r.Id == p.resp {
|
|
pending = append(pending, Pending{Item: item, Response: r})
|
|
}
|
|
}
|
|
}
|
|
return pending, nil
|
|
}
|
|
|
|
func (s *Store) MarkDelivered(ctx context.Context, responseIDs []string) error {
|
|
return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
|
|
at := now()
|
|
for _, id := range responseIDs {
|
|
_, err := tx.ExecContext(ctx, `UPDATE inbox_responses SET delivered_at = ? WHERE id = ?`, at, id)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
})
|
|
}
|
|
|
|
func (s *Store) MarkUndelivered(ctx context.Context, responseIDs []string) error {
|
|
return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
|
|
for _, id := range responseIDs {
|
|
_, err := tx.ExecContext(ctx, `UPDATE inbox_responses SET delivered_at = NULL WHERE id = ?`, id)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
})
|
|
}
|
|
|
|
type querier interface {
|
|
QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row
|
|
QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error)
|
|
}
|
|
|
|
func getItem(ctx context.Context, q querier, id string) (api.InboxItem, error) {
|
|
item := api.InboxItem{}
|
|
kind := ""
|
|
payload := ""
|
|
state := ""
|
|
createdAt := ""
|
|
resolvedBy := sql.NullString{}
|
|
resolvedAt := sql.NullString{}
|
|
|
|
err := q.QueryRowContext(ctx, `
|
|
SELECT id, project_id, task_id, agent_session_id, thread_id, kind, payload, state,
|
|
created_at, resolved_by, resolved_at
|
|
FROM inbox_items WHERE id = ?
|
|
`, id).Scan(&item.Id, &item.ProjectId, &item.TaskId, &item.AgentSessionId, &item.ThreadId, &kind,
|
|
&payload, &state, &createdAt, &resolvedBy, &resolvedAt)
|
|
if err != nil {
|
|
return item, err
|
|
}
|
|
|
|
item.Kind = api.InboxItemKind(kind)
|
|
item.State = api.InboxState(state)
|
|
err = json.Unmarshal([]byte(payload), &item.Payload)
|
|
if err != nil {
|
|
return item, err
|
|
}
|
|
item.CreatedAt, err = parseTime(createdAt)
|
|
if err != nil {
|
|
return item, err
|
|
}
|
|
item.ResolvedAt, err = parseTimePtr(resolvedAt)
|
|
if err != nil {
|
|
return item, err
|
|
}
|
|
if resolvedBy.Valid {
|
|
item.ResolvedBy = &resolvedBy.String
|
|
}
|
|
|
|
item.Responses, err = getResponses(ctx, q, id)
|
|
return item, err
|
|
}
|
|
|
|
func getResponses(ctx context.Context, q querier, itemID string) ([]api.InboxResponse, error) {
|
|
rows, err := q.QueryContext(ctx, `
|
|
SELECT id, item_id, user_login, type, options, text, client, exit_code, images, attachments, created_at, delivered_at
|
|
FROM inbox_responses WHERE item_id = ? ORDER BY created_at
|
|
`, itemID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
responses := []api.InboxResponse{}
|
|
for rows.Next() {
|
|
r := api.InboxResponse{}
|
|
typ := ""
|
|
options := sql.NullString{}
|
|
text := sql.NullString{}
|
|
client := ""
|
|
exitCode := sql.NullInt64{}
|
|
images := ""
|
|
attachments := ""
|
|
createdAt := ""
|
|
delivered := sql.NullString{}
|
|
|
|
err = rows.Scan(&r.Id, &r.ItemId, &r.User, &typ, &options, &text, &client, &exitCode, &images, &attachments, &createdAt, &delivered)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if client != "" {
|
|
r.Client = &client
|
|
}
|
|
if exitCode.Valid {
|
|
code := int(exitCode.Int64)
|
|
r.ExitCode = &code
|
|
}
|
|
|
|
r.Type = api.InboxResponseType(typ)
|
|
if options.Valid {
|
|
chosen := []string{}
|
|
err = json.Unmarshal([]byte(options.String), &chosen)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
r.Options = &chosen
|
|
}
|
|
if text.Valid {
|
|
r.Text = &text.String
|
|
}
|
|
err = json.Unmarshal([]byte(images), &r.Images)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
err = json.Unmarshal([]byte(attachments), &r.Attachments)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
r.CreatedAt, err = parseTime(createdAt)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
r.DeliveredAt, err = parseTimePtr(delivered)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
responses = append(responses, r)
|
|
}
|
|
return responses, rows.Err()
|
|
}
|
|
|
|
func emitItem(ctx context.Context, tx *sql.Tx, emit emitFunc, kind, id, actor string) error {
|
|
item, err := getItem(ctx, tx, id)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return emit(api.Event{
|
|
Kind: kind,
|
|
ProjectId: &item.ProjectId,
|
|
TaskId: &item.TaskId,
|
|
AgentSessionId: &item.AgentSessionId,
|
|
Actor: actor,
|
|
Payload: ToPayload(item),
|
|
})
|
|
}
|
|
|
|
func formatTimePtr(t *time.Time) any {
|
|
if t == nil {
|
|
return nil
|
|
}
|
|
return formatTime(*t)
|
|
}
|
|
|
|
func parseTimePtr(s sql.NullString) (*time.Time, error) {
|
|
if !s.Valid {
|
|
return nil, nil
|
|
}
|
|
t, err := parseTime(s.String)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &t, nil
|
|
}
|