Files
byre/internal/agents/manager.go
T

898 lines
21 KiB
Go

package agents
import (
"context"
"crypto/rand"
"encoding/json"
"errors"
"fmt"
"log/slog"
"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")
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
worktreeRoot string
mu sync.Mutex
live map[string]*live
closing bool
rateLimits map[api.AgentSessionVendor]map[string]any
merges map[string]string
deliverMu sync.Mutex
capMu sync.Mutex
snapMu sync.Mutex
snaps map[string]*snapState
bad *badLines
watchStop chan struct{}
}
type live struct {
sess *claude.Session
id string
dir string
model string
gitName string
gitEmail string
vendor api.AgentSessionVendor
taskID string
projectID string
stopping bool
interrupted bool
asked bool
bounced bool
compacting bool
rateLimited map[string]any
contextModel string
contextUsed int
contextWindow int
}
func New(st *store.Store, limits Limits, worktreeRoot string) *Manager {
return &Manager{
store: st,
limits: limits,
worktreeRoot: worktreeRoot,
live: map[string]*live{},
rateLimits: map[api.AgentSessionVendor]map[string]any{},
merges: map[string]string{},
snaps: map[string]*snapState{},
bad: newBadLines(),
watchStop: make(chan struct{}),
}
}
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()
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, 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 = m.store.SetSessionModel(ctx, sess.Id, model, actor)
if err != nil {
return api.AgentSession{}, err
}
sess.Model = model
}
task, err := m.store.GetTask(ctx, sess.TaskId)
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
}
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
}
err = m.retire(ctx, sess, actor)
if err != nil {
return err
}
}
return m.drain(ctx)
}
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()
}
if sess.State == api.AgentSessionStateQueued {
err := m.store.DeleteQueuedTurn(ctx, sess.Id)
if err != nil {
return err
}
}
return m.store.SetSessionState(ctx, sess.Id, api.AgentSessionStateCompleted, actor, 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)
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, nil)
if err != nil {
return err
}
slog.Info("resumed session", "session", sess.Id, "task", task.Id)
}
return m.drain(ctx)
}
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()
}
}
func (m *Manager) LiveCount() int {
m.mu.Lock()
defer m.mu.Unlock()
return len(m.live)
}
type Usage struct {
Running []api.TurnRef
Queued []api.TurnRef
Caps api.Caps
RateLimits map[string]map[string]any
}
func (m *Manager) Usage(ctx context.Context) (Usage, error) {
u := Usage{
Running: []api.TurnRef{},
Queued: []api.TurnRef{},
Caps: api.Caps{Global: m.limits.Global, PerVendor: map[string]int{}},
RateLimits: map[string]map[string]any{},
}
for vendor, n := range m.limits.PerVendor {
u.Caps.PerVendor[string(vendor)] = n
}
running, err := m.store.SessionsInState(ctx, api.AgentSessionStateRunning)
if err != nil {
return u, err
}
for _, sess := range running {
u.Running = append(u.Running, api.TurnRef{SessionId: sess.Id, TaskId: sess.TaskId, Vendor: string(sess.Vendor)})
}
queued, err := m.store.QueuedTurns(ctx)
if err != nil {
return u, err
}
for _, q := range queued {
u.Queued = append(u.Queued, api.TurnRef{SessionId: q.SessionID, TaskId: q.TaskID, Vendor: string(q.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) {
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, images, api.AgentSessionStateIdle)
}
err = m.store.EnqueueTurn(ctx, store.QueuedTurn{
SessionID: sess.Id,
Vendor: sess.Vendor,
Actor: actor,
Text: text,
Images: images,
CreatedAt: time.Now().UTC(),
})
if err != nil {
return api.AgentSession{}, err
}
err = transition(ctx, m.store, sess.Id, api.AgentSessionStateIdle, api.AgentSessionStateQueued, actor)
if err != nil {
derr := m.store.DeleteQueuedTurn(ctx, sess.Id)
if derr != nil {
slog.Error("remove queued turn", "session", sess.Id, "err", derr)
}
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, images []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 {
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.Control(cctx, map[string]any{"subtype": "set_model", "model": model})
if err != nil {
return fmt.Errorf("set model %s: %w", model, err)
}
l.model = model
return nil
}
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, q.Images, 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, images []string) error {
imgs := []claude.Image{}
for _, id := range images {
b, data, err := m.store.ReadBlob(ctx, id)
if err != nil {
return err
}
imgs = append(imgs, claude.Image{MediaType: b.ContentType, Data: data})
}
head, err := git.Head(l.dir)
if err != nil {
return err
}
payload := map[string]any{"text": text, "head": head}
if len(images) > 0 {
payload["images"] = images
}
err = m.emit(ctx, l, actor, "user.message", payload)
if err != nil {
return err
}
m.mu.Lock()
l.interrupted = false
l.asked = false
l.bounced = false
l.compacting = actor == "byre" && text == handoffPrompt
l.rateLimited = nil
m.mu.Unlock()
return l.sess.Send(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,
}
opts := claude.Options{
Cwd: task.WorktreePath,
Model: sess.Model,
GitName: creator.GitName,
GitEmail: creator.GitEmail,
Guidance: guidance,
MCPServers: []*mcp.Server{{Name: "byre", Version: "0.0.0", Tools: []mcp.Tool{m.askTool(l), m.cancelQuestionTool(l), m.createTaskTool(l), m.setMergeTitleTool(l)}}},
Hooks: map[string][]claude.HookMatcher{
"UserPromptSubmit": {{Func: remindGuidance}},
"Stop": {{Func: m.bounceProseQuestion(l)}},
"PostToolUse": {{Func: m.snapshotAfterTool(l)}},
},
OnBadLine: func(line []byte, reason string) { m.recordBadLine(l, line, reason) },
}
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)
go m.snapshotTimer(l)
return l, nil
}
func (m *Manager) pump(l *live) {
ctx := context.Background()
for ev := range l.sess.Events() {
start := time.Now()
if wait := start.Sub(ev.Read); wait > 200*time.Millisecond {
slog.Info("agent event waited", "session", l.id, "type", ev.Type, "wait", wait)
}
err := m.handle(ctx, l, ev)
if err != nil {
slog.Error("record agent event", "session", l.id, "type", ev.Type, "err", err)
}
if isDecodeError(err) {
m.recordBadLine(l, ev.Raw, fmt.Sprintf("decode %s: %s", ev.Type, err))
}
if took := time.Since(start); took > 200*time.Millisecond {
slog.Info("agent event slow", "session", l.id, "type", ev.Type, "took", took)
}
}
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
}
}
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
}
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
}
m.mu.Lock()
interrupted := l.interrupted
rateLimited := l.rateLimited
compacting := l.compacting
m.mu.Unlock()
finished := !res.IsError && !interrupted && strings.TrimSpace(res.Result) != ""
handoff := compacting && finished
payload := store.ToPayload(res)
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
}
m.snapshotSoon(l, api.SnapshotTriggerTurnEnd)
switch {
case handoff:
payload["handoff"] = true
case finished && !compacting:
payload["inbox_item_id"] = m.fileSummary(l, res.Result)
}
err = m.emit(ctx, l, agentActor, "turn.ended", payload)
if err != nil {
return err
}
window := res.ModelUsage[l.contextModel].ContextWindow
if window > 0 {
err = m.updateContext(ctx, l, l.contextUsed, 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
}
err = m.drain(ctx)
if err != nil {
slog.Error("start queued turns", "err", err)
}
if handoff {
go m.continueFromHandoff(l.id, res.Result)
return nil
}
mergeBy, merging := m.takeMerge(l.taskID)
if merging && finished {
go m.finishMerge(l, mergeBy)
return nil
}
err = m.Deliver(ctx, l.id)
if err != nil {
slog.Error("deliver inbox responses", "session", l.id, "err", err)
}
return nil
case "stream_event":
text, err := ev.TextDelta()
if err != nil || text == "" {
return err
}
m.store.Broadcast(api.Event{
Kind: "agent.delta",
ProjectId: &l.projectID,
TaskId: &l.taskID,
AgentSessionId: &l.id,
Actor: agentActor,
Payload: 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)
}
func commitMessage(res claude.Result) 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) 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:])
}