Run Claude sessions on tasks with event streaming and restart recovery
This commit is contained in:
Ian Gulliver
@@ -0,0 +1,451 @@
|
||||
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
|
||||
}
|
||||
|
||||
type live struct {
|
||||
sess *claude.Session
|
||||
gate *report.Gate
|
||||
id string
|
||||
taskID string
|
||||
projectID string
|
||||
stopping 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)
|
||||
}
|
||||
|
||||
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.SetSessionState(ctx, sess.Id, api.AgentSessionStateRunning, actor, nil)
|
||||
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
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
}}
|
||||
|
||||
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()}}},
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
return m.store.SetSessionState(ctx, l.id, api.AgentSessionStateIdle, agentActor, 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:])
|
||||
}
|
||||
Reference in New Issue
Block a user