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 Manager struct {
	store *store.Store

	mu      sync.Mutex
	live    map[string]*live
	closing bool

	deliverMu sync.Mutex
}

type live struct {
	sess        *claude.Session
	gate        *report.Gate
	id          string
	taskID      string
	projectID   string
	stopping    bool
	interrupted bool
}

func New(st *store.Store) *Manager {
	return &Manager{store: st, live: map[string]*live{}}
}

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, false)
}

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, true)
}

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
		}

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

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

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

	return nil
}

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, true)
		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 nil
}

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

func (m *Manager) turn(ctx context.Context, sess api.AgentSession, task api.Task, actor, text string, resume bool) (api.AgentSession, error) {
	l, err := m.ensureLive(ctx, sess, task, resume)
	if err != nil {
		return api.AgentSession{}, err
	}

	err = m.store.TransitionSession(ctx, sess.Id, api.AgentSessionStateIdle, api.AgentSessionStateRunning, actor)
	if errors.Is(err, store.ErrStateChanged) {
		return api.AgentSession{}, fmt.Errorf("%w: %w", ErrConflict, err)
	}
	if err != nil {
		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) 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
	m.mu.Unlock()

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

func (m *Manager) ensureLive(ctx context.Context, sess api.AgentSession, task api.Task, resume bool) (*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, 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}},
		},
	}
	if resume {
		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.")
}

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
		m.mu.Unlock()

		if res.IsError && !interrupted {
			m.fileError(l, fmt.Sprintf("The turn ended with an error (%s): %s", res.Subtype, res.Result))
		}

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

		go func() {
			err := m.Deliver(context.Background(), l.id)
			if err != nil {
				slog.Error("deliver inbox responses", "session", l.id, "err", err)
			}
		}()
		return nil

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