byre/internal/agents/queue.go

package agents

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

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

type sessionLock struct {
	mu   sync.Mutex
	refs int
}

func (m *Manager) lockDelivery(sessionID string) func() {
	m.deliverMu.Lock()
	sl := m.deliverLocks[sessionID]
	if sl == nil {
		sl = &sessionLock{}
		m.deliverLocks[sessionID] = sl
	}
	sl.refs++
	m.deliverMu.Unlock()

	sl.mu.Lock()
	return func() {
		sl.mu.Unlock()
		m.deliverMu.Lock()
		sl.refs--
		if sl.refs == 0 {
			delete(m.deliverLocks, sessionID)
		}
		m.deliverMu.Unlock()
	}
}

func (m *Manager) Message(ctx context.Context, sessionID, actor, text, model, effort, sendID, client string, images ...string) (api.AgentSession, bool, error) {
	unlock := m.lockDelivery(sessionID)
	defer unlock()

	prior, sent, err := m.priorSend(ctx, sendID)
	if err != nil || sent {
		return prior, false, err
	}

	sess, err := m.store.GetSession(ctx, sessionID)
	if err != nil {
		return api.AgentSession{}, false, err
	}
	if sess.State == api.AgentSessionStateIdle && isLeadActor(actor) {
		err = m.checkNotStopped(ctx, sess)
		if err != nil {
			return api.AgentSession{}, false, err
		}
	}
	landing, err := m.store.HeldForLanding(ctx, sess.TaskId, sess.Id)
	if err != nil {
		return api.AgentSession{}, false, err
	}
	held := sess.State == api.AgentSessionStateIdle && sess.Role == api.Implementer && (automatedActor(actor) && m.interrupted(sess.Id) || landing || m.isPaused(sess.Id))
	if !held && (sess.State != api.AgentSessionStateRunning || sess.Role != api.Implementer) {
		sess, err = m.Send(ctx, sessionID, actor, text, model, effort, sendID, client, images...)
		return sess, false, err
	}

	sess, err = m.changeModel(ctx, sess, model, effort, actor)
	if err != nil {
		return api.AgentSession{}, false, err
	}

	task, err := m.store.GetTask(ctx, sess.TaskId)
	if err != nil {
		return api.AgentSession{}, false, err
	}
	err = m.store.AddTaskInput(ctx, store.TaskInput{TaskID: sess.TaskId, Member: sess.Member, Actor: actor, Text: text, Images: images, SendID: sendID, Client: client}, task.ProjectId, sess.Id)
	if err != nil {
		return api.AgentSession{}, false, err
	}
	slog.Info("message queued", "session", sess.Id, "task", sess.TaskId, "actor", actor, "interrupted", held)
	return sess, true, nil
}

func (m *Manager) priorSend(ctx context.Context, sendID string) (api.AgentSession, bool, error) {
	if sendID == "" {
		return api.AgentSession{}, false, nil
	}
	sessionID, found, err := m.store.SendSession(ctx, sendID)
	if err != nil || !found || sessionID == "" {
		return api.AgentSession{}, found, err
	}
	sess, err := m.store.GetSession(ctx, sessionID)
	if err != nil {
		return api.AgentSession{}, false, err
	}
	slog.Info("send already accepted", "session", sess.Id, "send", sendID)
	return sess, true, nil
}

func automatedActor(actor string) bool {
	return actor == "byre" || actor == agentActor || actor == boardActor
}

func (m *Manager) interrupted(sessionID string) bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	l := m.live[sessionID]
	return l != nil && l.interrupted
}

func (m *Manager) Withdraw(ctx context.Context, taskID string, inputID int64, actor string) error {
	task, err := m.store.GetTask(ctx, taskID)
	if err != nil {
		return err
	}

	unlock := m.lockDelivery(queueLockKey(taskID))
	defer unlock()

	sessionID, err := m.store.WithdrawTaskInput(ctx, taskID, inputID, task.ProjectId, actor)
	if errors.Is(err, store.ErrNotFound) {
		return fmt.Errorf("%w: message %d is not queued", ErrConflict, inputID)
	}
	if err != nil {
		return err
	}
	slog.Info("queued message withdrawn", "session", sessionID, "task", taskID, "input", inputID, "actor", actor)
	return nil
}

func queueLockKey(taskID string) string {
	return "task:" + taskID
}

func (m *Manager) deliverQueued(ctx context.Context, sess api.AgentSession) (bool, error) {
	unlock := m.lockDelivery(queueLockKey(sess.TaskId))
	defer unlock()

	inputs, err := m.store.TaskInputs(ctx, sess.TaskId, sess.Member)
	if err != nil || len(inputs) == 0 {
		return false, err
	}
	if sess.Member != "" {
		started, err := m.store.SessionStarted(ctx, sess.Id)
		if err != nil || !started {
			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.SendID, in.Client, 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})
}

type queued struct {
	actor  string
	text   string
	sendID string
	client string
}

type heldImages struct {
	id    int64
	actor string
}

func takeBeforeImages[T any](inputs []T, ref func(T) (int64, string, []string)) ([]T, []int64, heldImages) {
	out := []T{}
	ids := []int64{}
	for _, in := range inputs {
		id, actor, images := ref(in)
		if len(images) > 0 {
			return out, ids, heldImages{id, actor}
		}
		out = append(out, in)
		ids = append(ids, id)
	}
	return out, ids, heldImages{}
}

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

	unlock := m.lockDelivery(queueLockKey(l.taskID))
	defer unlock()

	inputs, err := m.store.TaskInputs(ctx, l.taskID, l.member)
	if err != nil {
		slog.Error("list queued messages", "session", l.id, "err", err)
		return nil, heldImages{}
	}
	out, ids, held := takeBeforeImages(inputs, func(in store.TaskInput) (int64, string, []string) {
		return in.ID, in.Actor, in.Images
	})
	if len(ids) == 0 {
		return nil, held
	}
	err = m.store.DeleteTaskInputs(ctx, ids)
	if err != nil {
		slog.Error("delete queued messages", "session", l.id, "err", err)
		return nil, held
	}
	return out, held
}

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

	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, heldImages{}
	}
	messages := []store.LeadInput{}
	for _, in := range inputs {
		if in.Kind == store.LeadMessage {
			messages = append(messages, in)
		}
	}
	out, ids, held := takeBeforeImages(messages, func(in store.LeadInput) (int64, string, []string) {
		return in.ID, in.Actor, in.Images
	})
	if len(ids) == 0 {
		return nil, held
	}
	err = m.store.DeleteLeadInputs(ctx, ids)
	if err != nil {
		slog.Error("delete held lead messages", "session", l.id, "err", err)
		return nil, held
	}
	return out, held
}

func (m *Manager) queuedParts(ctx context.Context, l *live, msgs []queued, held heldImages) ([]string, int64) {
	parts := []string{}
	if len(msgs) > 0 {
		parts = append(parts, m.recordQueued(ctx, l, msgs))
	}
	hint := m.imagesHint(l, held)
	if hint == "" {
		return parts, 0
	}
	return append(parts, hint), held.id
}

func (m *Manager) imagesHint(l *live, held heldImages) string {
	if held.id == 0 {
		return ""
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	if l.hintedImages == held.id {
		return ""
	}
	l.hintedImages = held.id
	return fmt.Sprintf("%s sent you a message with images, which byre can only hand you as a new turn. Finish the step you're on, end this turn with a short summary, and don't start more work: they're waiting for you to read it.", held.actor)
}

func (m *Manager) recordQueued(ctx context.Context, l *live, msgs []queued) string {
	b := &strings.Builder{}
	for i, q := range msgs {
		if l.implementer() && isMergeInPrompt(q.actor, q.text) {
			m.takeMergeTurn(l)
		}
		payload := map[string]any{"text": q.text, "mid_turn": true}
		if isFindingsMessage(q.actor, q.text) {
			payload["findings"] = true
		}
		if q.sendID != "" {
			payload["send_id"] = q.sendID
		}
		if q.client != "" {
			payload["client"] = q.client
		}
		err := m.emit(ctx, l, q.actor, "user.message", payload)
		if err != nil {
			slog.Error("record mid-turn message", "session", l.id, "err", err)
		}
		if i > 0 {
			b.WriteString("\n\n")
		}
		if q.actor == boardActor {
			fmt.Fprintf(b, "While you were working, %s\n\nTake it into account now.", q.text)
			continue
		}
		if isLeadActor(q.actor) {
			fmt.Fprintf(b, "%s (not %s) sent you this message while you were working; take it into account now:\n\n%s", actorName(q.actor), leadOwner(q.actor), q.text)
			continue
		}
		if q.client != "" {
			fmt.Fprintf(b, "%s sent you this message from %s while you were working; take it into account now:\n\n%s", q.actor, m.clientName(q.client), q.text)
			continue
		}
		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()
}