byre/internal/agents/process.go

package agents

import (
	"context"
	"crypto/rand"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"strings"
	"time"

	"g.fc.run/ai/byre/internal/antigravity"
	"g.fc.run/ai/byre/internal/api"
	"g.fc.run/ai/byre/internal/claude"
	"g.fc.run/ai/byre/internal/codex"
	"g.fc.run/ai/byre/internal/mcp"
)

type image struct {
	mediaType string
	data      []byte
}

type agentProcess interface {
	send(ctx context.Context, text string, images []image) error
	interrupt(ctx context.Context) error
	setModel(ctx context.Context, model string) error
	kill()
	pid() int
	lastLine() time.Time
	exited() bool
	drive(m *Manager, l *live)
	discard()
}

func vendorOf(model string) api.AgentSessionVendor {
	switch {
	case strings.HasPrefix(model, "gpt-"):
		return api.Codex
	case strings.HasPrefix(model, "gemini-"):
		return api.Antigravity
	}
	return api.Claude
}

func initialVendorSessionID(vendor api.AgentSessionVendor) string {
	if vendor != api.Claude {
		return ""
	}
	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:])
}

func (m *Manager) changeModel(ctx context.Context, sess api.AgentSession, model, actor string) (api.AgentSession, error) {
	if model == "" || model == sess.Model {
		return sess, nil
	}
	if vendorOf(model) != sess.Vendor {
		return api.AgentSession{}, fmt.Errorf("%w: session %s is %s; %s is a %s model", ErrInvalid, sess.Id, sess.Vendor, model, vendorOf(model))
	}
	err := m.store.SetSessionModel(ctx, sess.Id, model, actor)
	if err != nil {
		return api.AgentSession{}, err
	}
	sess.Model = model
	return sess, nil
}

type claudeProcess struct {
	s *claude.Session
}

func (m *Manager) startClaude(ctx context.Context, sess api.AgentSession, l *live, g string, servers []*mcp.Server, stdio []mcp.Stdio, env []string) (agentProcess, error) {
	hooks := map[string][]claude.HookMatcher{
		"UserPromptSubmit": {{Func: remindGuidance(g)}},
	}
	if l.reviewID == "" {
		hooks["PostToolUse"] = []claude.HookMatcher{{Func: m.afterTool(l)}}
		if l.owner == "" {
			hooks["Stop"] = []claude.HookMatcher{{Func: m.bounceProseQuestion(l)}}
		}
	}

	opts := claude.Options{
		Cwd:        l.dir,
		Model:      sess.Model,
		GitName:    l.gitName,
		GitEmail:   l.gitEmail,
		Env:        env,
		Guidance:   g + l.notes,
		MCPServers: servers,
		StdioMCP:   stdio,
		Hooks:      hooks,
		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
	}
	switch {
	case started:
		opts.Resume = sess.VendorSessionId
	case sess.ForkedFrom != nil:
		opts.SessionID = sess.VendorSessionId
		opts.Resume = *sess.ForkedFrom
		opts.ResumeAt = *sess.ForkedAt
		opts.ForkSession = true
	default:
		opts.SessionID = sess.VendorSessionId
	}

	s, err := claude.Start(ctx, opts)
	if err != nil {
		return nil, fmt.Errorf("start claude: %w", err)
	}
	return claudeProcess{s: s}, nil
}

func (p claudeProcess) send(ctx context.Context, text string, images []image) error {
	imgs := []claude.Image{}
	for _, img := range images {
		imgs = append(imgs, claude.Image{MediaType: img.mediaType, Data: img.data})
	}
	return p.s.Send(text, imgs...)
}

func (p claudeProcess) interrupt(ctx context.Context) error {
	return p.s.Interrupt(ctx)
}

func (p claudeProcess) setModel(ctx context.Context, model string) error {
	_, err := p.s.Control(ctx, map[string]any{"subtype": "set_model", "model": model})
	return err
}

func (p claudeProcess) kill()               { p.s.Kill() }
func (p claudeProcess) pid() int            { return p.s.Pid() }
func (p claudeProcess) lastLine() time.Time { return p.s.LastLine() }
func (p claudeProcess) exited() bool        { return p.s.Exited() }

func (p claudeProcess) discard() {
	for range p.s.Events() {
	}
}

func (p claudeProcess) drive(m *Manager, l *live) {
	ctx := context.Background()
	for ev := range p.s.Events() {
		m.handleProcessEvent(l, ev.Read, ev.Type, ev.Raw, func() error {
			return m.handle(ctx, l, ev)
		})
	}
}

type codexProcess struct {
	s *codex.Session
}

func (m *Manager) startCodex(ctx context.Context, sess api.AgentSession, l *live, g string, servers []*mcp.Server, stdio []mcp.Stdio, env []string) (agentProcess, error) {
	err := os.MkdirAll(m.codexHome, 0o700)
	if err != nil {
		return nil, err
	}

	opts := codex.Options{
		Cwd:        l.dir,
		Model:      sess.Model,
		GitName:    l.gitName,
		GitEmail:   l.gitEmail,
		Env:        append(env, "CODEX_HOME="+m.codexHome),
		Guidance:   g + l.notes,
		Reminder:   g,
		MCPServers: servers,
		StdioMCP:   stdio,
		OnBadLine:  func(line []byte, reason string) { m.recordBadLine(l, line, reason) },
	}
	switch {
	case sess.VendorSessionId != "":
		opts.Resume = sess.VendorSessionId
	case sess.ForkedFrom != nil:
		opts.ForkFrom = *sess.ForkedFrom
		opts.ForkAt = *sess.ForkedAt
	}

	s, err := codex.Start(ctx, opts)
	if err != nil {
		return nil, fmt.Errorf("start codex: %w", err)
	}
	if sess.VendorSessionId == "" {
		err = m.store.SetVendorSessionID(ctx, sess.Id, s.ID())
		if err != nil {
			s.Kill()
			return nil, err
		}
	}
	return codexProcess{s: s}, nil
}

func (p codexProcess) send(ctx context.Context, text string, images []image) error {
	imgs := []codex.Image{}
	for _, img := range images {
		imgs = append(imgs, codex.Image{MediaType: img.mediaType, Data: img.data})
	}
	_, err := p.s.Send(ctx, text, imgs...)
	return err
}

func (p codexProcess) interrupt(ctx context.Context) error {
	return p.s.Interrupt(ctx)
}

func (p codexProcess) setModel(ctx context.Context, model string) error {
	p.s.SetModel(model)
	return nil
}

func (p codexProcess) kill()               { p.s.Kill() }
func (p codexProcess) pid() int            { return p.s.Pid() }
func (p codexProcess) lastLine() time.Time { return p.s.LastLine() }
func (p codexProcess) exited() bool        { return p.s.Exited() }

func (p codexProcess) discard() {
	for range p.s.Events() {
	}
}

func (p codexProcess) drive(m *Manager, l *live) {
	ctx := context.Background()
	for ev := range p.s.Events() {
		m.handleProcessEvent(l, ev.Read, ev.Method, ev.Params, func() error {
			return m.handleCodex(ctx, l, p.s, ev)
		})
	}
}

type antigravityProcess struct {
	s *antigravity.Session
}

func (m *Manager) startAntigravity(ctx context.Context, sess api.AgentSession, l *live, g string, servers []*mcp.Server, stdio []mcp.Stdio, env []string) (agentProcess, error) {
	opts := antigravity.Options{
		Dir:        filepath.Join(m.antigravityRoot, sess.Id),
		Cwd:        l.dir,
		Model:      sess.Model,
		Resume:     sess.VendorSessionId,
		GitName:    l.gitName,
		GitEmail:   l.gitEmail,
		Env:        env,
		Guidance:   g,
		Notes:      strings.TrimSpace(l.notes),
		Reminder:   g,
		MCPServers: servers,
		StdioMCP:   stdio,
		Hooks:      map[string]antigravity.HookFunc{"PreInvocation": m.agyBeforeModel(l)},
		OnBadLine:  func(line []byte, reason string) { m.recordBadLine(l, line, reason) },
	}

	s, err := antigravity.Start(ctx, opts)
	if err != nil {
		return nil, fmt.Errorf("start agy: %w", err)
	}
	if sess.VendorSessionId == "" {
		err = m.store.SetVendorSessionID(ctx, sess.Id, s.ID())
		if err != nil {
			s.Kill()
			return nil, err
		}
	}
	return antigravityProcess{s: s}, nil
}

func (p antigravityProcess) send(ctx context.Context, text string, images []image) error {
	imgs := []antigravity.Image{}
	for _, img := range images {
		imgs = append(imgs, antigravity.Image{MediaType: img.mediaType, Data: img.data})
	}
	return p.s.Send(text, imgs...)
}

func (p antigravityProcess) interrupt(ctx context.Context) error {
	return p.s.Interrupt()
}

func (p antigravityProcess) setModel(ctx context.Context, model string) error {
	return p.s.SetModel(model)
}

func (p antigravityProcess) kill()               { p.s.Kill() }
func (p antigravityProcess) pid() int            { return p.s.Pid() }
func (p antigravityProcess) lastLine() time.Time { return p.s.LastLine() }
func (p antigravityProcess) exited() bool        { return p.s.Exited() }

func (p antigravityProcess) discard() {
	for range p.s.Events() {
	}
}

func (p antigravityProcess) drive(m *Manager, l *live) {
	ctx := context.Background()
	steps := &agySteps{text: map[int]*strings.Builder{}, tools: map[string]bool{}}
	for ev := range p.s.Events() {
		m.handleProcessEvent(l, ev.Read, ev.Kind, ev.Raw, func() error {
			return m.handleAntigravity(ctx, l, p, steps, ev)
		})
	}
}

func (m *Manager) handleProcessEvent(l *live, read time.Time, name string, raw []byte, handle func() error) {
	start := time.Now()
	if wait := start.Sub(read); wait > 200*time.Millisecond {
		slog.Info("agent event waited", "session", l.id, "event", name, "wait", wait)
	}
	err := handle()
	if err != nil {
		slog.Error("record agent event", "session", l.id, "event", name, "err", err)
	}
	if isDecodeError(err) {
		m.recordBadLine(l, raw, fmt.Sprintf("decode %s: %s", name, err))
	}
	if took := time.Since(start); took > 200*time.Millisecond {
		slog.Info("agent event slow", "session", l.id, "event", name, "took", took)
	}
}

func (m *Manager) continueProseQuestion(ctx context.Context, l *live, p agentProcess, final string, success bool) (bool, error) {
	if !success || !l.implementer() {
		return false, nil
	}
	bounce, err := m.takeBounce(ctx, l, final)
	if err != nil || !bounce {
		return false, err
	}
	m.mu.Lock()
	l.lastText = ""
	m.mu.Unlock()
	go func() {
		err := p.send(context.Background(), proseQuestionNudge, nil)
		if err != nil {
			slog.Error("send prose question nudge", "session", l.id, "err", err)
		}
	}()
	return true, nil
}