byre/internal/agents/antigravity.go

package agents

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"strings"
	"sync"
	"time"

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

type agySteps struct {
	mu    sync.Mutex
	text  map[int]*strings.Builder
	tools map[string]bool
}

func newAgySteps() *agySteps {
	return &agySteps{text: map[int]*strings.Builder{}, tools: map[string]bool{}}
}

func (m *Manager) handleAntigravity(ctx context.Context, l *live, p antigravityProcess, ev antigravity.Event) error {
	switch ev.Kind {
	case "init":
		return nil

	case "step_update":
		su, err := ev.StepUpdate()
		if err != nil {
			return err
		}
		return m.agyStep(ctx, l, p, ev, su)

	case "result":
		res, err := ev.Result()
		if err != nil {
			return err
		}
		return m.agyTurnEnded(ctx, l, p, res)
	}

	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) agyStep(ctx context.Context, l *live, p antigravityProcess, ev antigravity.Event, su antigravity.StepUpdate) error {
	switch su.StepType {
	case "user_input":
		return nil

	case "agent_response":
		p.steps.mu.Lock()
		b := p.steps.text[su.StepIndex]
		if b == nil {
			b = &strings.Builder{}
			p.steps.text[su.StepIndex] = b
		}
		b.WriteString(su.TextDelta)
		done := su.State != "ACTIVE"
		if done {
			delete(p.steps.text, su.StepIndex)
		}
		p.steps.mu.Unlock()

		if su.TextDelta != "" {
			m.store.Broadcast(m.event(l, agentActor, "agent.delta", map[string]any{"text": su.TextDelta}))
		}
		text := strings.TrimSpace(b.String())
		if !done || text == "" {
			return nil
		}
		m.mu.Lock()
		l.lastText = text
		m.mu.Unlock()
		return m.emit(ctx, l, agentActor, "agent.message", map[string]any{"text": text})

	case "tool":
		id := fmt.Sprintf("%s/%d", su.ConversationID, su.StepIndex)
		info := antigravity.ToolInfo{Name: su.ToolName}
		if su.ToolInfo != nil {
			info = *su.ToolInfo
		}
		p.steps.mu.Lock()
		called := p.steps.tools[id]
		p.steps.tools[id] = su.State == "ACTIVE"
		if su.State != "ACTIVE" {
			delete(p.steps.tools, id)
		}
		p.steps.mu.Unlock()

		if !called {
			name, input, err := agyToolCall(su.ToolName, info.Parameters)
			if err != nil {
				return err
			}
			err = m.emit(ctx, l, agentActor, "agent.tool_call", map[string]any{"id": id, "name": name, "input": input})
			if err != nil {
				return err
			}
		}
		if su.State == "ACTIVE" {
			return nil
		}
		content := info.Output
		if info.Error != "" {
			content = info.Error
		}
		return m.emit(ctx, l, agentActor, "agent.tool_result", map[string]any{"tool_use_id": id, "is_error": su.State != "DONE" || info.Error != "", "content": content})
	}

	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 agyToolCall(name string, params json.RawMessage) (string, any, error) {
	switch name {
	case "run_command":
		p := struct {
			CommandLine string
		}{}
		err := json.Unmarshal(params, &p)
		return "Bash", map[string]any{"command": p.CommandLine}, err
	case "call_mcp_tool":
		p := struct {
			ServerName string
			ToolName   string
			Arguments  json.RawMessage
		}{}
		err := json.Unmarshal(params, &p)
		if err != nil {
			return "", nil, err
		}
		args, err := decode(p.Arguments)
		return "mcp__" + strings.TrimPrefix(p.ServerName, "byre_") + "__" + p.ToolName, args, err
	}
	input, err := decode(params)
	return name, input, err
}

func (m *Manager) agyTurnEnded(ctx context.Context, l *live, p antigravityProcess, res antigravity.Result) error {
	m.mu.Lock()
	final := l.lastText
	m.mu.Unlock()

	success := res.Status == "SUCCESS"
	bounce := false
	if success && l.implementer() {
		var err error
		bounce, err = m.takeBounce(ctx, l, final)
		if err != nil {
			return err
		}
	}
	if bounce {
		err := m.emit(ctx, l, "byre", "user.message", map[string]any{"text": proseQuestionNudge, "mid_turn": true, "bounce": true})
		if err != nil {
			return err
		}
		m.mu.Lock()
		l.lastText = ""
		m.mu.Unlock()
		go func() {
			err := p.s.Send(proseQuestionNudge)
			if err != nil {
				slog.Error("send prose question nudge", "session", l.id, "err", err)
			}
		}()
		return nil
	}

	r := endResult{
		subtype: strings.ToLower(res.Status),
		isError: !success,
		result:  final,
	}
	if !success {
		r.result = res.Error
	}
	r.payload = map[string]any{
		"subtype":  r.subtype,
		"is_error": r.isError,
		"result":   r.result,
	}
	return m.turnEnded(ctx, l, r)
}

var agyWindowMins = map[string]int{"5h": 300, "weekly": 7 * 24 * 60}

func (m *Manager) readAgyRateLimit() {
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	groups, err := antigravity.ReadUsage(ctx)
	if err != nil {
		slog.Error("read antigravity rate limits", "err", err)
		return
	}
	m.setVendorUsage(api.Antigravity, agyUsage(groups))
}

func agyUsage(groups []antigravity.UsageGroup) api.Usage {
	u := api.Usage{UpdatedAt: time.Now(), Windows: []api.UsageWindow{}}
	for _, g := range groups {
		for _, b := range g.Buckets {
			u.Windows = append(u.Windows, percentWindow(b.ID, 1-b.RemainingFraction, b.ResetTime.Unix(), int64(agyWindowMins[b.Window])*60))
		}
	}
	return u
}

func (m *Manager) agyBeforeModel(l *live) antigravity.HookFunc {
	return func(ctx context.Context, input json.RawMessage) (map[string]any, error) {
		in := struct {
			InvocationNum int `json:"invocationNum"`
		}{}
		err := json.Unmarshal(input, &in)
		if err != nil {
			return nil, err
		}
		if in.InvocationNum == 0 || l.reviewID != "" {
			return map[string]any{}, nil
		}
		mt := m.deliverMidTurn(ctx, l)
		if mt.text == "" {
			return map[string]any{}, nil
		}
		return map[string]any{"injectSteps": []map[string]any{{"userMessage": mt.text}}}, nil
	}
}