byre/internal/agents/manager.go

package agents

import (
	"context"
	"crypto/rand"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

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

var (
	ErrConflict      = errors.New("conflict")
	ErrInvalid       = errors.New("invalid")
	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 Manager struct {
	store        *store.Store
	worktreeRoot string
	leadRoot     string
	codexHome    string

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

	deliverMu sync.Mutex
	leadMu    sync.Mutex
	signals   chan api.Event

	snapMu sync.Mutex
	snaps  map[string]*snapState

	mergeMu      sync.Mutex
	mergeHolders map[string]string

	bad       *badLines
	watchStop chan struct{}
}

type live struct {
	sess        agentProcess
	id          string
	dir         string
	model       string
	gitName     string
	gitEmail    string
	vendor      api.AgentSessionVendor
	taskID      string
	projectID   string
	owner       string
	stopping    bool
	interrupted bool
	asked       bool
	askedBlock  bool
	bounced     bool
	compacting  bool
	wantCompact bool
	rateLimited map[string]any
	lastText    string

	contextModel  string
	contextUsed   int
	contextWindow int
}

func New(st *store.Store, worktreeRoot, leadRoot, codexHome string) *Manager {
	m := &Manager{
		store:        st,
		worktreeRoot: worktreeRoot,
		leadRoot:     leadRoot,
		codexHome:    codexHome,
		live:         map[string]*live{},
		rateLimits:   map[api.AgentSessionVendor]map[string]any{},
		leadAnswered: map[string]bool{},
		paneFocused:  map[string]bool{},
		signals:      make(chan api.Event, 1024),
		snaps:        map[string]*snapState{},
		mergeHolders: map[string]string{},
		bad:          newBadLines(),
		watchStop:    make(chan struct{}),
	}
	go m.signalLoop()
	return m
}

func (m *Manager) taskFor(ctx context.Context, sess api.AgentSession) (api.Task, error) {
	if sess.Owner != "" {
		return api.Task{}, nil
	}
	return m.store.GetTask(ctx, sess.TaskId)
}

func (m *Manager) Create(ctx context.Context, task api.Task, creator api.User, prompt, model string, images ...string) (api.AgentSession, error) {
	return m.create(ctx, task, creator, prompt, model, creator.Login, images)
}

func (m *Manager) create(ctx context.Context, task api.Task, creator api.User, prompt, model, actor string, images []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()
	vendor := vendorOf(model)
	sess := api.AgentSession{
		Id:              id,
		TaskId:          task.Id,
		Vendor:          vendor,
		Model:           model,
		VendorSessionId: initialVendorSessionID(vendor, 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, actor, prompt, images)
}

func (m *Manager) Send(ctx context.Context, sessionID, actor, text, model string, images ...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)
	}

	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
	}

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

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

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
	}

	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
		}
		err = m.retire(ctx, sess, actor)
		if err != nil {
			return err
		}
	}

	return nil
}

func (m *Manager) retire(ctx context.Context, sess api.AgentSession, actor string) error {
	m.mu.Lock()
	l := m.live[sess.Id]
	if l != nil {
		l.stopping = true
	}
	m.mu.Unlock()

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

	return m.store.SetSessionState(ctx, sess.Id, api.AgentSessionStateCompleted, actor, nil)
}

func (m *Manager) Recover(ctx context.Context) error {
	err := m.recoverMerges(ctx)
	if err != nil {
		return err
	}

	sessions, err := m.store.SessionsInState(ctx, api.AgentSessionStateRunning)
	if err != nil {
		return err
	}

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

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

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

	idle, err := m.store.SessionsInState(ctx, api.AgentSessionStateIdle)
	if err != nil {
		return err
	}
	for _, sess := range idle {
		m.logUndelivered(ctx, sess.Id, "idle at startup")
	}

	go m.deliverAllLeads()
	go m.startHeldMerges(context.Background())
	return nil
}

func (m *Manager) logUndelivered(ctx context.Context, sessionID, reason string) {
	pending, err := m.store.Undelivered(ctx, sessionID)
	if err != nil {
		slog.Error("list undelivered responses", "session", sessionID, "err", err)
		return
	}
	if len(pending) == 0 {
		return
	}
	slog.Info("undelivered responses", "session", sessionID, "reason", reason, "responses", responseIDs(pending))
}

func (m *Manager) Close() {
	close(m.watchStop)

	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()
	}

	err := m.parkLeads(context.Background())
	if err != nil {
		slog.Error("park lead sessions", "err", err)
	}
}

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

type Usage struct {
	Running    []api.TurnRef
	RateLimits map[string]map[string]any
}

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

	running, err := m.store.SessionsInState(ctx, api.AgentSessionStateRunning)
	if err != nil {
		return u, err
	}
	for _, sess := range running {
		if sess.Owner != "" {
			continue
		}
		u.Running = append(u.Running, api.TurnRef{SessionId: sess.Id, TaskId: sess.TaskId, Vendor: string(sess.Vendor)})
	}

	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, images []string) (api.AgentSession, error) {
	err := transition(ctx, m.store, sess.Id, api.AgentSessionStateIdle, api.AgentSessionStateRunning, actor)
	if err != nil {
		return api.AgentSession{}, err
	}

	l, err := m.ensureLive(ctx, sess, task)
	if err == nil {
		err = m.syncModel(ctx, l, sess.Model)
	}
	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, images)
	if err != nil {
		return api.AgentSession{}, err
	}

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

func (m *Manager) syncModel(ctx context.Context, l *live, model string) error {
	if model == l.model {
		return nil
	}

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

	err := l.sess.setModel(cctx, model)
	if err != nil {
		return fmt.Errorf("set model %s: %w", model, err)
	}
	l.model = model
	return nil
}

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, images []string) error {
	imgs := []image{}
	for _, id := range images {
		b, data, err := m.store.ReadBlob(ctx, id)
		if err != nil {
			return err
		}
		imgs = append(imgs, image{mediaType: b.ContentType, data: data})
	}

	payload := map[string]any{"text": text}
	if l.owner == "" {
		head, err := git.Head(l.dir)
		if err != nil {
			return err
		}
		payload["head"] = head
	}
	if len(images) > 0 {
		payload["images"] = images
	}
	compacting := isHandoffPrompt(actor, text)
	if compacting {
		payload["handoff"] = true
	}
	if isContinuation(actor, text) {
		payload["continuation"] = true
	}
	err := m.emit(ctx, l, actor, "user.message", payload)
	if err != nil {
		return err
	}

	m.mu.Lock()
	l.interrupted = false
	l.asked = false
	l.askedBlock = false
	l.bounced = false
	l.wantCompact = false
	l.compacting = compacting
	l.rateLimited = nil
	l.lastText = ""
	m.mu.Unlock()

	return l.sess.send(ctx, text, imgs)
}

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,
		dir:       task.WorktreePath,
		model:     sess.Model,
		gitName:   creator.GitName,
		gitEmail:  creator.GitEmail,
		vendor:    sess.Vendor,
		taskID:    task.Id,
		projectID: task.ProjectId,
		owner:     sess.Owner,
	}

	g, err := m.taskGuidance(ctx, task)
	if err != nil {
		return nil, err
	}
	tools := []mcp.Tool{m.askTool(l), m.cancelQuestionTool(l), m.createTaskTool(l), m.setMergeTitleTool(l), m.exposePortTool(l)}
	env := portEnv(task)
	if l.owner != "" {
		l.dir = filepath.Join(m.leadRoot, url.PathEscape(l.owner))
		err = os.MkdirAll(l.dir, 0o755)
		if err != nil {
			return nil, err
		}
		g = leadGuidance
		tools = m.leadTools(l)
		env = nil
	}
	servers := []*mcp.Server{{Name: "byre", Version: "0.0.0", Tools: tools}}

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

	switch l.vendor {
	case api.Codex:
		l.sess, err = m.startCodex(sctx, sess, l, g, servers, env)
	default:
		l.sess, err = m.startClaude(sctx, sess, l, g, servers, env)
	}
	if err != nil {
		return nil, err
	}

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

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

	l.sess.drive(m, l)

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

	if quiet {
		m.logUndelivered(ctx, l.id, "agent process ended while closing or stopping")
		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
	}

	m.logUndelivered(ctx, l.id, "agent process exited mid-turn")
	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)
	}
	if l.owner == "" {
		m.abandonMerge(l.taskID, l.projectID)
		m.fileError(l, "The agent process exited mid-turn. Send a message to resume it.")
	}
}

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
		}
		parent, err := ev.ParentToolUseID()
		if err != nil {
			return err
		}
		uuid, err := ev.UUID()
		if err != nil {
			return err
		}
		for _, b := range blocks {
			err = m.handleAssistantBlock(ctx, l, b, parent, uuid)
			if err != nil {
				return err
			}
		}
		if parent != "" {
			return nil
		}
		usage, err := ev.ContextUsage()
		if err != nil {
			return err
		}
		if usage.Tokens == 0 {
			return nil
		}
		l.contextModel = usage.Model
		return m.updateContext(ctx, l, usage.Tokens, l.contextWindow)

	case "user":
		blocks, err := ev.Blocks()
		if err != nil {
			return err
		}
		parent, err := ev.ParentToolUseID()
		if err != nil {
			return err
		}
		uuid, err := ev.UUID()
		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", withParent(map[string]any{
				"tool_use_id": b.ToolUseID,
				"is_error":    b.IsError,
				"content":     content,
			}, parent, uuid))
			if err != nil {
				return err
			}
		}
		return nil

	case "result":
		res, err := ev.Result()
		if err != nil {
			return err
		}
		return m.turnEnded(ctx, l, endResult{
			subtype: res.Subtype,
			isError: res.IsError,
			result:  res.Result,
			payload: store.ToPayload(res),
			window:  res.ModelUsage[l.contextModel].ContextWindow,
		})

	case "stream_event":
		text, err := ev.TextDelta()
		if err != nil || text == "" {
			return err
		}
		m.store.Broadcast(m.event(l, agentActor, "agent.delta", map[string]any{"text": text}))
		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)
}

type endResult struct {
	subtype string
	isError bool
	result  string
	payload map[string]any
	window  int
}

func (m *Manager) turnEnded(ctx context.Context, l *live, res endResult) error {
	if l.owner != "" {
		return m.leadTurnEnded(ctx, l, res)
	}

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

	finished := !res.isError && !interrupted && strings.TrimSpace(res.result) != ""
	handoff := compacting && finished

	payload := res.payload
	sha, err := git.CommitAll(l.dir, commitMessage(res), l.gitName, l.gitEmail)
	if err != nil {
		slog.Error("commit turn", "session", l.id, "err", err)
	}
	if sha != "" {
		payload["commit"] = sha
	}
	switch {
	case handoff:
		payload["handoff"] = true
	case finished && !compacting && !askedBlock:
		payload["inbox_item_id"] = m.fileSummary(l, res.result)
	}
	err = m.emit(ctx, l, agentActor, "turn.ended", payload)
	if err != nil {
		return err
	}
	m.snapshotSoon(l, api.SnapshotTriggerTurnEnd, 0)

	if res.window > 0 {
		err = m.updateContext(ctx, l, l.contextUsed, res.window)
		if err != nil {
			return err
		}
	}

	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
	}

	if handoff {
		m.logUndelivered(ctx, l.id, "turn ended in handoff")
		go m.continueFromHandoff(l.id, res.result)
		return nil
	}
	mergeBy := m.takeMerge(l.taskID)
	switch {
	case mergeBy != "" && finished:
		m.logUndelivered(ctx, l.id, "turn ended in merge")
		go m.finishMerge(l, mergeBy)
		return nil
	case mergeBy != "":
		go m.releaseMerge(l.projectID, l.taskID)
	default:
		go m.nextMerge(l.projectID)
	}
	err = m.Deliver(ctx, l.id)
	if err != nil {
		slog.Error("deliver inbox responses", "session", l.id, "err", err)
	}
	return nil
}

func commitMessage(res endResult) string {
	line, _, _ := strings.Cut(strings.TrimSpace(res.result), "\n")
	sentence, _, found := strings.Cut(line, ". ")
	if found {
		sentence += "."
	}
	sentence = strings.TrimSpace(strings.ReplaceAll(sentence, "**", ""))
	if sentence == "" {
		return fmt.Sprintf("Work from a turn that ended without a summary (%s)", res.subtype)
	}
	runes := []rune(sentence)
	if len(runes) > 100 {
		return string(runes[:99]) + "…"
	}
	return sentence
}

func (m *Manager) updateContext(ctx context.Context, l *live, used, window int) error {
	if used == l.contextUsed && window == l.contextWindow {
		return nil
	}
	l.contextUsed = used
	l.contextWindow = window
	return m.emit(ctx, l, agentActor, "session.context", map[string]any{"model": l.contextModel, "used": used, "window": window})
}

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

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

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

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

func withParent(payload map[string]any, parent, uuid string) map[string]any {
	if parent != "" {
		payload["parent_tool_use_id"] = parent
		return payload
	}
	if uuid != "" {
		payload["vendor_uuid"] = uuid
	}
	return payload
}

func (m *Manager) emit(ctx context.Context, l *live, actor, kind string, payload map[string]any) error {
	return m.store.AppendEvent(ctx, m.event(l, actor, kind, payload))
}

func (m *Manager) event(l *live, actor, kind string, payload map[string]any) api.Event {
	ev := api.Event{
		Kind:           kind,
		AgentSessionId: &l.id,
		Actor:          actor,
		Payload:        payload,
	}
	if l.owner != "" {
		ev.Owner = &l.owner
		return ev
	}
	ev.ProjectId = &l.projectID
	ev.TaskId = &l.taskID
	return ev
}

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:])
}