byre/internal/agents/process.go

package agents

import (
	"context"
	"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"
	"g.fc.run/ai/byre/internal/workflow"
)

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

func checkModel(sess api.AgentSession, model string) error {
	if workflow.VendorOf(model) != sess.Vendor {
		return fmt.Errorf("%w: session %s is %s; %s is a %s model", ErrInvalid, sess.Id, sess.Vendor, model, workflow.VendorOf(model))
	}
	return 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)}},
		"Stop":             {{Func: m.bounceProseQuestion(l)}},
		"PostToolUse":      {{Func: m.afterTool(l)}},
	}
	if l.owner != "" {
		hooks = map[string][]claude.HookMatcher{
			"UserPromptSubmit": {{Func: remindGuidance(g)}},
			"PostToolUse":      {{Func: m.afterTool(l)}},
		}
	}
	if l.reviewID != "" {
		hooks = map[string][]claude.HookMatcher{"UserPromptSubmit": {{Func: remindGuidance(g)}}}
	}

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

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() {
		start := time.Now()
		if wait := start.Sub(ev.Read); wait > 200*time.Millisecond {
			slog.Info("agent event waited", "session", l.id, "method", ev.Method, "wait", wait)
		}
		err := m.handleCodex(ctx, l, p.s, ev)
		if err != nil {
			slog.Error("record agent event", "session", l.id, "method", ev.Method, "err", err)
		}
		if isDecodeError(err) {
			m.recordBadLine(l, ev.Params, fmt.Sprintf("decode %s: %s", ev.Method, err))
		}
		if took := time.Since(start); took > 200*time.Millisecond {
			slog.Info("agent event slow", "session", l.id, "method", ev.Method, "took", took)
		}
	}
}

type antigravityProcess struct {
	s     *antigravity.Session
	steps *agySteps
}

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, steps: newAgySteps()}, 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()
	for ev := range p.s.Events() {
		start := time.Now()
		if wait := start.Sub(ev.Read); wait > 200*time.Millisecond {
			slog.Info("agent event waited", "session", l.id, "event", ev.Kind, "wait", wait)
		}
		err := m.handleAntigravity(ctx, l, p, ev)
		if err != nil {
			slog.Error("record agent event", "session", l.id, "event", ev.Kind, "err", err)
		}
		if isDecodeError(err) {
			m.recordBadLine(l, ev.Raw, fmt.Sprintf("decode %s: %s", ev.Kind, err))
		}
		if took := time.Since(start); took > 200*time.Millisecond {
			slog.Info("agent event slow", "session", l.id, "event", ev.Kind, "took", took)
		}
	}
}