byre/internal/agents/queue.go

package agents

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"strings"

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

func (m *Manager) Message(ctx context.Context, sessionID, actor, text, model string, images ...string) (api.AgentSession, error) {
	m.deliverMu.Lock()
	defer m.deliverMu.Unlock()

	sess, err := m.store.GetSession(ctx, sessionID)
	if err != nil {
		return api.AgentSession{}, err
	}
	if sess.State != api.AgentSessionStateRunning || sess.Role != api.Implementer {
		return m.Send(ctx, sessionID, actor, text, model, images...)
	}

	if model != "" && model != sess.Model {
		err = checkModel(sess, model)
		if err != nil {
			return api.AgentSession{}, err
		}
		err = m.store.SetSessionModel(ctx, sess.Id, model, actor)
		if err != nil {
			return api.AgentSession{}, err
		}
		sess.Model = model
	}

	err = m.store.AddTaskInput(ctx, store.TaskInput{TaskID: sess.TaskId, Actor: actor, Text: text, Images: images})
	if err != nil {
		return api.AgentSession{}, err
	}
	slog.Info("message queued", "session", sess.Id, "task", sess.TaskId, "actor", actor)
	return sess, nil
}

func (m *Manager) deliverQueued(ctx context.Context, sess api.AgentSession) (bool, error) {
	inputs, err := m.store.TaskInputs(ctx, sess.TaskId)
	if err != nil || len(inputs) == 0 {
		return false, err
	}

	task, err := m.store.GetTask(ctx, sess.TaskId)
	if err != nil {
		return false, err
	}

	in := inputs[0]
	_, err = m.Send(ctx, sess.Id, in.Actor, in.Text, "", in.Images...)
	if errors.Is(err, ErrConflict) {
		slog.Info("queued message held: session changed state", "session", sess.Id, "task", task.Id, "err", err)
		return true, nil
	}
	if err != nil {
		return false, err
	}
	slog.Info("queued message sent", "session", sess.Id, "task", task.Id, "input", in.ID)
	return true, m.store.DeleteTaskInputs(ctx, []int64{in.ID})
}

func (m *Manager) deliverAllQueued(ctx context.Context) {
	idle, err := m.store.SessionsInState(ctx, api.AgentSessionStateIdle)
	if err != nil {
		slog.Error("list idle sessions", "err", err)
		return
	}
	for _, sess := range idle {
		if sess.Role != api.Implementer {
			continue
		}
		m.deliverMu.Lock()
		_, err = m.deliverQueued(ctx, sess)
		m.deliverMu.Unlock()
		if err != nil {
			slog.Error("deliver queued message", "session", sess.Id, "err", err)
		}
	}
}

type queued struct {
	actor string
	text  string
}

func (m *Manager) takeQueuedMidTurn(ctx context.Context, l *live) []queued {
	m.mu.Lock()
	compacting := l.compacting
	m.mu.Unlock()
	if compacting {
		return nil
	}

	inputs, err := m.store.TaskInputs(ctx, l.taskID)
	if err != nil {
		slog.Error("list queued messages", "session", l.id, "err", err)
		return nil
	}
	out := []queued{}
	ids := []int64{}
	for _, in := range inputs {
		if len(in.Images) > 0 {
			break
		}
		out = append(out, queued{actor: in.Actor, text: in.Text})
		ids = append(ids, in.ID)
	}
	if len(ids) == 0 {
		return nil
	}
	err = m.store.DeleteTaskInputs(ctx, ids)
	if err != nil {
		slog.Error("delete queued messages", "session", l.id, "err", err)
		return nil
	}
	return out
}

func (m *Manager) takeLeadQueuedMidTurn(ctx context.Context, l *live) []queued {
	m.mu.Lock()
	compacting := l.compacting
	m.mu.Unlock()
	if compacting {
		return nil
	}

	m.leadMu.Lock()
	defer m.leadMu.Unlock()

	inputs, err := m.store.LeadInputs(ctx, l.owner)
	if err != nil {
		slog.Error("list held lead messages", "session", l.id, "err", err)
		return nil
	}
	out := []queued{}
	ids := []int64{}
	for _, in := range inputs {
		if in.Kind != store.LeadMessage {
			continue
		}
		if len(in.Images) > 0 {
			break
		}
		out = append(out, queued{actor: in.Actor, text: in.Text})
		ids = append(ids, in.ID)
	}
	if len(ids) == 0 {
		return nil
	}
	err = m.store.DeleteLeadInputs(ctx, ids)
	if err != nil {
		slog.Error("delete held lead messages", "session", l.id, "err", err)
		return nil
	}
	return out
}

func (m *Manager) recordQueued(ctx context.Context, l *live, msgs []queued) string {
	b := &strings.Builder{}
	for i, q := range msgs {
		err := m.emit(ctx, l, q.actor, "user.message", map[string]any{"text": q.text, "mid_turn": true})
		if err != nil {
			slog.Error("record mid-turn message", "session", l.id, "err", err)
		}
		if i > 0 {
			b.WriteString("\n\n")
		}
		fmt.Fprintf(b, "%s sent you this message while you were working; take it into account now:\n\n%s", q.actor, q.text)
	}
	slog.Info("queued messages sent mid-turn", "session", l.id, "count", len(msgs))
	return b.String()
}