Add Antigravity (agy) as a third agent backend for gemini-* models
This commit is contained in:
Ian Gulliver
@@ -0,0 +1,581 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user