byre/internal/agents/manager.go

package agents

import (
	"context"
	"crypto/rand"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"sync"
	"time"

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

var (
	ErrConflict      = errors.New("conflict")
	ErrNoGitIdentity = errors.New("set your git name and email (PUT /api/me) before starting agents")
)

const agentActor = "agent"

const restartNotice = "byre restarted while you were mid-turn, so your last turn was cut off. Re-check the working tree (git status and the files you were editing) before continuing where you left off."

type Limits struct {
	Global    int
	PerVendor map[api.AgentSessionVendor]int
}

var DefaultLimits = Limits{
	Global:    8,
	PerVendor: map[api.AgentSessionVendor]int{api.Claude: 6},
}

type Manager struct {
	store  *store.Store
	limits Limits

	mu         sync.Mutex
	live       map[string]*live
	closing    bool
	rateLimits map[api.AgentSessionVendor]map[string]any

	deliverMu sync.Mutex
	capMu     sync.Mutex
}

type live struct {
	sess        *claude.Session
	gate        *report.Gate
	id          string
	vendor      api.AgentSessionVendor
	taskID      string
	projectID   string
	stopping    bool
	interrupted bool
	rateLimited map[string]any
}

func New(st *store.Store, limits Limits) *Manager {
	return &Manager{
		store:      st,
		limits:     limits,
		live:       map[string]*live{},
		rateLimits: map[api.AgentSessionVendor]map[string]any{},
	}
}

func (m *Manager) Create(ctx context.Context, task api.Task, creator api.User, prompt, model string) (api.AgentSession, error) {
	if creator.GitName == "" || creator.GitEmail == "" {
		return api.AgentSession{}, ErrNoGitIdentity
	}

	if task.Status != api.TaskStatusActive {
		return api.AgentSession{}, fmt.Errorf("%w: task %s is %s", ErrConflict, task.Id, task.Status)
	}

	existing, err := m.store.ListSessions(ctx, task.Id)
	if err != nil {
		return api.AgentSession{}, err
	}
	for _, s := range existing {
		if s.State != api.AgentSessionStateCompleted {
			return api.AgentSession{}, fmt.Errorf("%w: task %s already has session %s", ErrConflict, task.Id, s.Id)
		}
	}

	id := newUUID()
	sess := api.AgentSession{
		Id:              id,
		TaskId:          task.Id,
		Vendor:          api.Claude,
		Model:           model,
		VendorSessionId: id,
		State:           api.AgentSessionStateIdle,
		Role:            api.Implementer,
		CreatedBy:       creator.Login,
		CreatedAt:       time.Now().UTC(),
	}

	err = m.store.CreateSession(ctx, sess, task.ProjectId)
	if err != nil {
		return api.AgentSession{}, err
	}

	return m.turn(ctx, sess, task, creator.Login, prompt)
}

func (m *Manager) Send(ctx context.Context, sessionID, actor, text string) (api.AgentSession, error) {
	sess, err := m.store.GetSession(ctx, sessionID)
	if err != nil {
		return api.AgentSession{}, err
	}

	if sess.State != api.AgentSessionStateIdle {
		return api.AgentSession{}, fmt.Errorf("%w: session %s is %s", ErrConflict, sess.Id, sess.State)
	}

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

	return m.turn(ctx, sess, task, actor, text)
}

func (m *Manager) Interrupt(ctx context.Context, sessionID string) (api.AgentSession, error) {
	sess, err := m.store.GetSession(ctx, sessionID)
	if err != nil {
		return api.AgentSession{}, err
	}

	if sess.State == api.AgentSessionStateQueued {
		return m.cancelQueued(ctx, sess)
	}

	m.mu.Lock()
	l := m.live[sessionID]
	m.mu.Unlock()

	if sess.State != api.AgentSessionStateRunning || l == nil {
		return api.AgentSession{}, fmt.Errorf("%w: session %s has no turn running", ErrConflict, sess.Id)
	}

	m.mu.Lock()
	l.interrupted = true
	m.mu.Unlock()

	ictx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()

	err = l.sess.Interrupt(ictx)
	if err != nil {
		return api.AgentSession{}, err
	}

	return sess, nil
}

func (m *Manager) StopTask(ctx context.Context, taskID, actor string) error {
	sessions, err := m.store.ListSessions(ctx, taskID)
	if err != nil {
		return err
	}

	for _, sess := range sessions {
		if sess.State == api.AgentSessionStateCompleted {
			continue
		}

		m.mu.Lock()
		l := m.live[sess.Id]
		if l != nil {
			l.stopping = true
		}
		m.mu.Unlock()

		if l != nil {
			l.sess.Kill()
		}

		if sess.State == api.AgentSessionStateQueued {
			err = m.store.DeleteQueuedTurn(ctx, sess.Id)
			if err != nil {
				return err
			}
		}

		err = m.store.SetSessionState(ctx, sess.Id, api.AgentSessionStateCompleted, actor, nil)
		if err != nil {
			return err
		}
	}

	return m.drain(ctx)
}

func (m *Manager) Recover(ctx context.Context) error {
	sessions, err := m.store.SessionsInState(ctx, api.AgentSessionStateRunning)
	if err != nil {
		return err
	}

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

		l, err := m.ensureLive(ctx, sess, task)
		if err != nil {
			slog.Error("resume session", "session", sess.Id, "err", err)
			err = m.store.SetSessionState(ctx, sess.Id, api.AgentSessionStateIdle, "byre", map[string]any{"error": err.Error()})
			if err != nil {
				return err
			}
			continue
		}

		err = m.message(ctx, l, "byre", restartNotice)
		if err != nil {
			return err
		}
		slog.Info("resumed session", "session", sess.Id, "task", task.Id)
	}

	return m.drain(ctx)
}

func (m *Manager) Close() {
	m.mu.Lock()
	m.closing = true
	all := []*live{}
	for _, l := range m.live {
		all = append(all, l)
	}
	m.mu.Unlock()

	for _, l := range all {
		l.sess.Kill()
	}
}

func (m *Manager) LiveCount() int {
	m.mu.Lock()
	defer m.mu.Unlock()
	return len(m.live)
}

type Usage struct {
	Running    int
	Queued     int
	RateLimits map[string]map[string]any
}

func (m *Manager) Usage(ctx context.Context) (Usage, error) {
	u := Usage{RateLimits: map[string]map[string]any{}}

	counts, err := m.store.CountRunning(ctx)
	if err != nil {
		return u, err
	}
	for _, n := range counts {
		u.Running += n
	}

	queued, err := m.store.QueuedTurns(ctx)
	if err != nil {
		return u, err
	}
	u.Queued = len(queued)

	m.mu.Lock()
	defer m.mu.Unlock()
	for vendor, info := range m.rateLimits {
		u.RateLimits[string(vendor)] = info
	}
	return u, nil
}

func (m *Manager) turn(ctx context.Context, sess api.AgentSession, task api.Task, actor, text string) (api.AgentSession, error) {
	m.capMu.Lock()
	defer m.capMu.Unlock()

	room, err := m.hasRoom(ctx, sess.Vendor)
	if err != nil {
		return api.AgentSession{}, err
	}

	if room {
		return m.start(ctx, sess, task, actor, text, api.AgentSessionStateIdle)
	}

	err = transition(ctx, m.store, sess.Id, api.AgentSessionStateIdle, api.AgentSessionStateQueued, actor)
	if err != nil {
		return api.AgentSession{}, err
	}

	err = m.store.EnqueueTurn(ctx, store.QueuedTurn{
		SessionID: sess.Id,
		Vendor:    sess.Vendor,
		Actor:     actor,
		Text:      text,
		CreatedAt: time.Now().UTC(),
	})
	if err != nil {
		return api.AgentSession{}, err
	}

	return m.store.GetSession(ctx, sess.Id)
}

func (m *Manager) start(ctx context.Context, sess api.AgentSession, task api.Task, actor, text string, from api.AgentSessionState) (api.AgentSession, error) {
	err := transition(ctx, m.store, sess.Id, from, api.AgentSessionStateRunning, actor)
	if err != nil {
		return api.AgentSession{}, err
	}

	l, err := m.ensureLive(ctx, sess, task)
	if err != nil {
		serr := m.store.SetSessionState(ctx, sess.Id, api.AgentSessionStateIdle, actor, map[string]any{"error": err.Error()})
		if serr != nil {
			slog.Error("record failed start", "session", sess.Id, "err", serr)
		}
		return api.AgentSession{}, err
	}

	err = m.message(ctx, l, actor, text)
	if err != nil {
		return api.AgentSession{}, err
	}

	return m.store.GetSession(ctx, sess.Id)
}

func (m *Manager) hasRoom(ctx context.Context, vendor api.AgentSessionVendor) (bool, error) {
	counts, err := m.store.CountRunning(ctx)
	if err != nil {
		return false, err
	}

	total := 0
	for _, n := range counts {
		total += n
	}

	return total < m.limits.Global && counts[vendor] < m.limits.PerVendor[vendor], nil
}

func (m *Manager) drain(ctx context.Context) error {
	m.capMu.Lock()
	defer m.capMu.Unlock()

	queued, err := m.store.QueuedTurns(ctx)
	if err != nil {
		return err
	}

	for _, q := range queued {
		room, err := m.hasRoom(ctx, q.Vendor)
		if err != nil {
			return err
		}
		if !room {
			continue
		}

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

		err = m.store.DeleteQueuedTurn(ctx, q.SessionID)
		if err != nil {
			return err
		}

		_, err = m.start(ctx, sess, task, q.Actor, q.Text, api.AgentSessionStateQueued)
		if err != nil {
			slog.Error("start queued turn", "session", q.SessionID, "err", err)
		}
	}

	return nil
}

func (m *Manager) cancelQueued(ctx context.Context, sess api.AgentSession) (api.AgentSession, error) {
	m.capMu.Lock()
	defer m.capMu.Unlock()

	err := m.store.DeleteQueuedTurn(ctx, sess.Id)
	if errors.Is(err, store.ErrNotFound) {
		return api.AgentSession{}, fmt.Errorf("%w: session %s has no queued turn", ErrConflict, sess.Id)
	}
	if err != nil {
		return api.AgentSession{}, err
	}

	err = transition(ctx, m.store, sess.Id, api.AgentSessionStateQueued, api.AgentSessionStateIdle, "byre")
	if err != nil {
		return api.AgentSession{}, err
	}

	return m.store.GetSession(ctx, sess.Id)
}

func transition(ctx context.Context, st *store.Store, id string, from, to api.AgentSessionState, actor string) error {
	err := st.TransitionSession(ctx, id, from, to, actor)
	if errors.Is(err, store.ErrStateChanged) {
		return fmt.Errorf("%w: %w", ErrConflict, err)
	}
	return err
}

func (m *Manager) message(ctx context.Context, l *live, actor, text string) error {
	err := m.emit(ctx, l, actor, "user.message", map[string]any{"text": text})
	if err != nil {
		return err
	}

	m.mu.Lock()
	l.interrupted = false
	l.rateLimited = nil
	m.mu.Unlock()

	l.gate.Reset()
	return l.sess.Send(text)
}

func (m *Manager) ensureLive(ctx context.Context, sess api.AgentSession, task api.Task) (*live, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	if m.closing {
		return nil, errors.New("shutting down")
	}

	l := m.live[sess.Id]
	if l != nil {
		return l, nil
	}

	creator, err := m.store.GetUser(ctx, sess.CreatedBy)
	if err != nil {
		return nil, err
	}

	l = &live{id: sess.Id, vendor: sess.Vendor, taskID: task.Id, projectID: task.ProjectId}
	l.gate = &report.Gate{OnReport: func(r report.Report) {
		err := m.emit(context.Background(), l, agentActor, "agent.report", store.ToPayload(r))
		if err != nil {
			slog.Error("record report", "session", l.id, "err", err)
		}
		m.fileReport(l, r)
	}}

	opts := claude.Options{
		Cwd:        task.WorktreePath,
		Model:      sess.Model,
		GitName:    creator.GitName,
		GitEmail:   creator.GitEmail,
		MCPServers: []*mcp.Server{{Name: "byre", Version: "0.0.0", Tools: []mcp.Tool{l.gate.Tool(), m.askTool(l)}}},
		Hooks: map[string][]claude.HookMatcher{
			"Stop": {{Func: l.gate.StopHook}},
		},
	}
	started, err := m.store.SessionStarted(ctx, sess.Id)
	if err != nil {
		return nil, err
	}
	if started {
		opts.Resume = sess.VendorSessionId
	} else {
		opts.SessionID = sess.VendorSessionId
	}

	sctx, cancel := context.WithTimeout(ctx, 60*time.Second)
	defer cancel()

	l.sess, err = claude.Start(sctx, opts)
	if err != nil {
		return nil, fmt.Errorf("start claude: %w", err)
	}

	m.live[sess.Id] = l
	go m.pump(l)
	return l, nil
}

func (m *Manager) pump(l *live) {
	ctx := context.Background()

	for ev := range l.sess.Events() {
		err := m.handle(ctx, l, ev)
		if err != nil {
			slog.Error("record agent event", "session", l.id, "type", ev.Type, "err", err)
		}
	}

	m.mu.Lock()
	delete(m.live, l.id)
	quiet := m.closing || l.stopping
	m.mu.Unlock()

	if quiet {
		return
	}

	sess, err := m.store.GetSession(ctx, l.id)
	if err != nil {
		slog.Error("load exited session", "session", l.id, "err", err)
		return
	}

	if sess.State != api.AgentSessionStateRunning {
		return
	}

	err = m.store.SetSessionState(ctx, l.id, api.AgentSessionStateIdle, agentActor, map[string]any{"error": "agent process exited mid-turn"})
	if err != nil {
		slog.Error("record session exit", "session", l.id, "err", err)
	}
	m.fileError(l, "The agent process exited mid-turn. Send a message to resume it.")

	err = m.drain(ctx)
	if err != nil {
		slog.Error("start queued turns", "err", err)
	}
}

func (m *Manager) handle(ctx context.Context, l *live, ev claude.Event) error {
	switch ev.Type {
	case "assistant":
		blocks, err := ev.Blocks()
		if err != nil {
			return err
		}
		for _, b := range blocks {
			err = m.handleAssistantBlock(ctx, l, b)
			if err != nil {
				return err
			}
		}
		return nil

	case "user":
		blocks, err := ev.Blocks()
		if err != nil {
			return err
		}
		for _, b := range blocks {
			if b.Type != "tool_result" {
				continue
			}
			content, err := decode(b.Content)
			if err != nil {
				return err
			}
			err = m.emit(ctx, l, agentActor, "agent.tool_result", map[string]any{
				"tool_use_id": b.ToolUseID,
				"is_error":    b.IsError,
				"content":     content,
			})
			if err != nil {
				return err
			}
		}
		return nil

	case "result":
		res, err := ev.Result()
		if err != nil {
			return err
		}
		err = m.emit(ctx, l, agentActor, "turn.ended", store.ToPayload(res))
		if err != nil {
			return err
		}

		m.mu.Lock()
		interrupted := l.interrupted
		rateLimited := l.rateLimited
		m.mu.Unlock()

		failure := fmt.Sprintf("The turn ended with an error (%s): %s", res.Subtype, res.Result)
		switch {
		case !res.IsError || interrupted:
		case rateLimited != nil:
			m.fileRateLimited(l, rateLimited, failure)
		default:
			m.fileError(l, failure)
		}

		err = m.store.SetSessionState(ctx, l.id, api.AgentSessionStateIdle, agentActor, nil)
		if err != nil {
			return err
		}

		err = m.drain(ctx)
		if err != nil {
			slog.Error("start queued turns", "err", err)
		}
		err = m.Deliver(ctx, l.id)
		if err != nil {
			slog.Error("deliver inbox responses", "session", l.id, "err", err)
		}
		return nil

	case "rate_limit_event":
		info := struct {
			RateLimitInfo map[string]any `json:"rate_limit_info"`
		}{}
		err := json.Unmarshal(ev.Raw, &info)
		if err != nil {
			return err
		}

		m.mu.Lock()
		m.rateLimits[l.vendor] = info.RateLimitInfo
		if info.RateLimitInfo["status"] == "rejected" {
			l.rateLimited = info.RateLimitInfo
		}
		m.mu.Unlock()

	case "system":
		if ev.Subtype == "init" {
			return nil
		}
	}

	raw := map[string]any{}
	err := json.Unmarshal(ev.Raw, &raw)
	if err != nil {
		return err
	}
	return m.emit(ctx, l, agentActor, "agent.vendor", raw)
}

func (m *Manager) handleAssistantBlock(ctx context.Context, l *live, b claude.Block) error {
	switch b.Type {
	case "text":
		return m.emit(ctx, l, agentActor, "agent.message", map[string]any{"text": b.Text})

	case "thinking":
		if b.Thinking == "" {
			return nil
		}
		return m.emit(ctx, l, agentActor, "agent.thinking", map[string]any{"text": b.Thinking})

	case "tool_use":
		input, err := decode(b.Input)
		if err != nil {
			return err
		}
		return m.emit(ctx, l, agentActor, "agent.tool_call", map[string]any{
			"id":    b.ID,
			"name":  b.Name,
			"input": input,
		})
	}

	return m.emit(ctx, l, agentActor, "agent.vendor", store.ToPayload(b))
}

func (m *Manager) emit(ctx context.Context, l *live, actor, kind string, payload map[string]any) error {
	return m.store.AppendEvent(ctx, api.Event{
		Kind:           kind,
		ProjectId:      &l.projectID,
		TaskId:         &l.taskID,
		AgentSessionId: &l.id,
		Actor:          actor,
		Payload:        payload,
	})
}

func decode(raw json.RawMessage) (any, error) {
	if len(raw) == 0 {
		return nil, nil
	}
	var v any
	err := json.Unmarshal(raw, &v)
	return v, err
}

func newUUID() string {
	b := make([]byte, 16)
	rand.Read(b)
	b[6] = b[6]&0x0f | 0x40
	b[8] = b[8]&0x3f | 0x80
	return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:])
}