Add Antigravity (agy) as a third agent backend for gemini-* models

internal/agents/antigravity.go

@@ -0,0 +1,215 @@
+package agents
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"log/slog"
+	"strings"
+	"sync"
+
+	"g.fc.run/ai/byre/internal/antigravity"
+)
+
+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)
+}
+
+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
+	}
+}