byre/internal/store/inbox.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
}