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 string, images ...string) (api.AgentSession, bool, error) {
	unlock := m.lockDelivery(sessionID)
	defer unlock()

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

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

	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}, 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 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.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
		}
		unlock := m.lockDelivery(sess.Id)
		_, err = m.deliverQueued(ctx, sess)
		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) []store.TaskInput {
	m.mu.Lock()
	compacting := l.compacting
	m.mu.Unlock()
	if compacting {
		return nil
	}

	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
	}
	out := []store.TaskInput{}
	ids := []int64{}
	for _, in := range inputs {
		if len(in.Images) > 0 {
			break
		}
		out = append(out, in)
		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) []store.LeadInput {
	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 := []store.LeadInput{}
	ids := []int64{}
	for _, in := range inputs {
		if in.Kind != store.LeadMessage {
			continue
		}
		if len(in.Images) > 0 {
			break
		}
		out = append(out, in)
		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 {
		if l.implementer() && isMergeInPrompt(q.actor, q.text) {
			m.takeMergeTurn(l)
		}
		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")
		}
		if q.actor == boardActor {
			fmt.Fprintf(b, "While you were working, %s\n\nTake it into account now.", 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()
}