byre/internal/antigravity/session.go

package antigravity

import (
	"bufio"
	"bytes"
	"context"
	"crypto/rand"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"g.fc.run/ai/byre/internal/mcp"
)

const maxRuleBytes = 24000

type HookFunc func(ctx context.Context, input json.RawMessage) (map[string]any, error)

type Options struct {
	Dir        string
	Cwd        string
	Model      string
	Resume     string
	GitName    string
	GitEmail   string
	Env        []string
	Guidance   string
	Notes      string
	Reminder   string
	MCPServers []*mcp.Server
	StdioMCP   []mcp.Stdio
	Hooks      map[string]HookFunc
	OnBadLine  func(line []byte, reason string)
}

type Image struct {
	MediaType string
	Data      []byte
}

type Event struct {
	Kind string
	Raw  json.RawMessage
	Read time.Time
}

type Usage struct {
	InputTokens     int `json:"input_tokens"`
	OutputTokens    int `json:"output_tokens"`
	ThinkingTokens  int `json:"thinking_tokens"`
	CacheReadTokens int `json:"cache_read_tokens"`
	TotalTokens     int `json:"total_tokens"`
}

type ToolInfo struct {
	Name       string          `json:"name"`
	Parameters json.RawMessage `json:"parameters"`
	Output     string          `json:"output"`
	Error      string          `json:"error"`
}

type StepUpdate struct {
	ConversationID  string    `json:"conversation_id"`
	StepIndex       int       `json:"step_index"`
	State           string    `json:"state"`
	StepType        string    `json:"step_type"`
	TextDelta       string    `json:"text_delta"`
	ToolName        string    `json:"tool_name"`
	ToolInfo        *ToolInfo `json:"tool_info"`
	DurationSeconds float64   `json:"duration_seconds"`
	Usage           *Usage    `json:"usage"`
}

type Result struct {
	ConversationID  string  `json:"conversation_id"`
	Status          string  `json:"status"`
	Response        string  `json:"response"`
	Error           string  `json:"error"`
	DurationSeconds float64 `json:"duration_seconds"`
	NumTurns        int     `json:"num_turns"`
	Usage           Usage   `json:"usage"`
}

var knownEvents = map[string]bool{"init": true, "step_update": true, "result": true}

func (e Event) payload(v any) error {
	fields := map[string]json.RawMessage{}
	err := json.Unmarshal(e.Raw, &fields)
	if err != nil {
		return err
	}
	return json.Unmarshal(fields[e.Kind], v)
}

func (e Event) StepUpdate() (StepUpdate, error) {
	su := StepUpdate{}
	return su, e.payload(&su)
}

func (e Event) Result() (Result, error) {
	res := Result{}
	return res, e.payload(&res)
}

type proc struct {
	cmd      *exec.Cmd
	stdin    io.WriteCloser
	done     chan struct{}
	expected bool
}

type Session struct {
	opts Options

	mu     sync.Mutex
	p      *proc
	model  string
	killed bool
	images int

	idMu   sync.Mutex
	convID string

	events   chan Event
	dead     chan struct{}
	dieOnce  sync.Once
	started  chan struct{}
	initOnce sync.Once

	srv      *http.Server
	base     string
	servers  map[string]*mcp.Server
	lastLine atomic.Int64
}

func Start(ctx context.Context, opts Options) (*Session, error) {
	if len(opts.Guidance) > maxRuleBytes || len(opts.Notes) > maxRuleBytes {
		return nil, fmt.Errorf("guidance (%d bytes) and notes (%d bytes) must each fit agy's %d-byte rule limit", len(opts.Guidance), len(opts.Notes), maxRuleBytes)
	}

	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return nil, err
	}
	token := rand.Text()

	s := &Session{
		opts:    opts,
		model:   opts.Model,
		convID:  opts.Resume,
		events:  make(chan Event, 256),
		dead:    make(chan struct{}),
		started: make(chan struct{}),
		base:    "http://" + ln.Addr().String() + "/" + token,
		servers: map[string]*mcp.Server{},
	}
	for _, srv := range opts.MCPServers {
		s.servers[srv.Name] = srv
	}
	s.lastLine.Store(time.Now().UnixNano())

	mux := http.NewServeMux()
	mux.HandleFunc("POST /"+token+"/mcp/{server}", s.serveMCP)
	mux.HandleFunc("POST /"+token+"/hook/{event}", s.serveHook)
	s.srv = &http.Server{Handler: mux}
	go s.srv.Serve(ln)

	err = s.writePlugin()
	if err != nil {
		s.die()
		return nil, err
	}

	s.mu.Lock()
	err = s.spawnLocked()
	s.mu.Unlock()
	if err != nil {
		s.die()
		return nil, err
	}

	select {
	case <-s.started:
		return s, nil
	case <-s.dead:
		return nil, errors.New("agy exited before it started a conversation")
	case <-ctx.Done():
		s.Kill()
		return nil, ctx.Err()
	}
}

func (s *Session) pluginDir() string {
	return filepath.Join(s.opts.Dir, ".agents", "plugins", "byre")
}

func (s *Session) writePlugin() error {
	dir := s.pluginDir()
	err := os.MkdirAll(filepath.Join(dir, "rules"), 0o755)
	if err != nil {
		return err
	}

	servers := map[string]any{}
	for _, srv := range s.opts.MCPServers {
		servers[srv.Name] = map[string]any{"url": s.base + "/mcp/" + srv.Name}
	}
	for _, srv := range s.opts.StdioMCP {
		servers[srv.Name] = map[string]any{"command": srv.Command, "args": srv.Args}
	}

	hooks := map[string]any{}
	for event := range s.opts.Hooks {
		hooks[event] = []map[string]any{{
			"type":    "command",
			"command": fmt.Sprintf("curl -sS --max-time 25 --data-binary @- %s/hook/%s", s.base, event),
			"timeout": 30,
		}}
	}

	files := map[string]any{
		"plugin.json":     map[string]any{"name": "byre"},
		"mcp_config.json": map[string]any{"mcpServers": servers},
		"hooks.json":      map[string]any{"byre": hooks},
	}
	for name, v := range files {
		b, err := json.MarshalIndent(v, "", "  ")
		if err != nil {
			return err
		}
		err = os.WriteFile(filepath.Join(dir, name), b, 0o644)
		if err != nil {
			return err
		}
	}

	err = os.WriteFile(filepath.Join(dir, "rules", "AGENTS.md"), []byte(s.opts.Guidance), 0o644)
	if err != nil {
		return err
	}
	if s.opts.Notes == "" {
		return nil
	}
	return os.WriteFile(filepath.Join(dir, "rules", "notes.md"), []byte("---\ntrigger: always_on\n---\n"+s.opts.Notes), 0o644)
}

func (s *Session) serveMCP(w http.ResponseWriter, r *http.Request) {
	srv := s.servers[r.PathValue("server")]
	if srv == nil {
		http.NotFound(w, r)
		return
	}
	body, err := io.ReadAll(r.Body)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	msg := struct {
		ID json.RawMessage `json:"id"`
	}{}
	err = json.Unmarshal(body, &msg)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	if len(msg.ID) == 0 {
		w.WriteHeader(http.StatusAccepted)
		return
	}
	resp, err := srv.Handle(r.Context(), body)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	w.Write(resp)
}

func (s *Session) serveHook(w http.ResponseWriter, r *http.Request) {
	event := r.PathValue("event")
	fn := s.opts.Hooks[event]
	if fn == nil {
		http.NotFound(w, r)
		return
	}
	body, err := io.ReadAll(r.Body)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	out, err := fn(r.Context(), body)
	if err != nil {
		slog.Error("agy hook failed", "event", event, "err", err)
		out = map[string]any{}
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(out)
}

func (s *Session) spawnLocked() error {
	args := []string{
		"--input-format", "stream-json",
		"--output-format", "stream-json",
		"--dangerously-skip-permissions",
		"--add-dir", s.opts.Dir,
	}
	if s.model != "" {
		args = append(args, "--model", s.model)
	}
	if id := s.ID(); id != "" {
		args = append(args, "--conversation", id)
	}
	args = append(args, "--print", "")

	cmd := exec.Command("agy", args...)
	cmd.Dir = s.opts.Cwd
	cmd.Stderr = os.Stderr
	cmd.Env = append(os.Environ(),
		"GIT_AUTHOR_NAME="+s.opts.GitName,
		"GIT_AUTHOR_EMAIL="+s.opts.GitEmail,
		"GIT_COMMITTER_NAME="+s.opts.GitName,
		"GIT_COMMITTER_EMAIL="+s.opts.GitEmail,
	)
	cmd.Env = append(cmd.Env, s.opts.Env...)

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return err
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	err = cmd.Start()
	if err != nil {
		return fmt.Errorf("start agy: %w", err)
	}

	p := &proc{cmd: cmd, stdin: stdin, done: make(chan struct{})}
	s.p = p
	go s.readLoop(p, stdout)
	return nil
}

func (s *Session) readLoop(p *proc, stdout io.Reader) {
	r := bufio.NewReader(stdout)
	for {
		line, err := r.ReadBytes('\n')
		line = bytes.TrimRight(line, "\r\n")
		if len(line) > 0 {
			s.lastLine.Store(time.Now().UnixNano())
			s.handleLine(line)
		}
		if err != nil {
			if err != io.EOF {
				slog.Error("read agy stdout", "err", err)
			}
			break
		}
	}

	p.cmd.Wait()
	close(p.done)

	s.mu.Lock()
	expected := p.expected && !s.killed
	s.mu.Unlock()
	if !expected {
		s.die()
	}
}

func (s *Session) handleLine(line []byte) {
	read := time.Now()
	if line[0] != '{' {
		s.badLine(line, "not a JSON object")
		return
	}
	env := struct {
		Event          string `json:"event"`
		ConversationID string `json:"conversation_id"`
	}{}
	err := json.Unmarshal(line, &env)
	if err != nil {
		s.badLine(line, err.Error())
		return
	}
	if !knownEvents[env.Event] {
		s.badLine(line, fmt.Sprintf("unrecognised event %q", env.Event))
	}
	if env.Event == "init" {
		s.idMu.Lock()
		if s.convID == "" {
			s.convID = env.ConversationID
		}
		s.idMu.Unlock()
		s.initOnce.Do(func() { close(s.started) })
	}
	s.events <- Event{Kind: env.Event, Raw: append(json.RawMessage{}, line...), Read: read}
}

func (s *Session) die() {
	s.dieOnce.Do(func() {
		close(s.events)
		close(s.dead)
		s.srv.Close()
		err := os.RemoveAll(s.opts.Dir)
		if err != nil {
			slog.Error("remove agy session dir", "dir", s.opts.Dir, "err", err)
		}
	})
}

func (s *Session) restartLocked() error {
	p := s.p
	select {
	case <-p.done:
	case <-time.After(10 * time.Second):
		p.cmd.Process.Kill()
		<-p.done
	}
	return s.spawnLocked()
}

func (s *Session) Events() <-chan Event {
	return s.events
}

func (s *Session) ID() string {
	s.idMu.Lock()
	defer s.idMu.Unlock()
	return s.convID
}

func (s *Session) Pid() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.p.cmd.Process.Pid
}

func (s *Session) LastLine() time.Time {
	return time.Unix(0, s.lastLine.Load())
}

func (s *Session) Exited() bool {
	select {
	case <-s.dead:
		return true
	default:
		return false
	}
}

type block struct {
	Type string `json:"type"`
	Text string `json:"text"`
}

func (s *Session) Send(text string, images ...Image) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.killed {
		return errors.New("agy session was killed")
	}
	if s.p.expected {
		err := s.restartLocked()
		if err != nil {
			s.die()
			return err
		}
	}

	content := []block{}
	if s.opts.Reminder != "" {
		content = append(content, block{Type: "text", Text: "<system-reminder>\n" + s.opts.Reminder + "\n</system-reminder>"})
	}
	for _, img := range images {
		s.images++
		path := filepath.Join(s.opts.Dir, "images", fmt.Sprintf("%d.%s", s.images, strings.TrimPrefix(img.MediaType, "image/")))
		err := os.MkdirAll(filepath.Dir(path), 0o755)
		if err != nil {
			return err
		}
		err = os.WriteFile(path, img.Data, 0o644)
		if err != nil {
			return err
		}
		content = append(content, block{Type: "text", Text: "The human attached an image, saved at " + path + "; view it with view_file."})
	}
	content = append(content, block{Type: "text", Text: text})

	b, err := json.Marshal(map[string]any{
		"event":   "user",
		"message": map[string]any{"role": "user", "content": content},
	})
	if err != nil {
		return err
	}
	_, err = s.p.stdin.Write(append(b, '\n'))
	return err
}

func (s *Session) Interrupt() error {
	s.mu.Lock()
	p := s.p
	p.expected = true
	err := p.cmd.Process.Signal(syscall.SIGINT)
	s.mu.Unlock()
	if err != nil {
		return err
	}

	go func() {
		<-p.done
		s.mu.Lock()
		defer s.mu.Unlock()
		if s.killed {
			s.die()
			return
		}
		if s.p != p {
			return
		}
		err := s.spawnLocked()
		if err != nil {
			slog.Error("restart agy after interrupt", "err", err)
			s.die()
		}
	}()
	return nil
}

func (s *Session) SetModel(model string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if model == s.model {
		return nil
	}
	s.model = model
	s.p.expected = true
	s.p.stdin.Close()
	err := s.restartLocked()
	if err != nil {
		s.die()
	}
	return err
}

func (s *Session) Kill() {
	s.mu.Lock()
	s.killed = true
	p := s.p
	s.mu.Unlock()

	if p != nil {
		p.cmd.Process.Kill()
		<-p.done
	}
	s.die()
}

func (s *Session) badLine(line []byte, reason string) {
	slog.Warn("bad agy line", "reason", reason, "line", string(line))
	if s.opts.OnBadLine != nil {
		s.opts.OnBadLine(line, reason)
	}
}