476 lines
12 KiB
Go
476 lines
12 KiB
Go
package antigravity
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"regexp"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"github.com/chenhg5/cc-connect/core"
|
|
)
|
|
|
|
// antigravitySession manages multi-turn conversations with the Antigravity CLI (agy).
|
|
type antigravitySession struct {
|
|
cmd string
|
|
workDir string
|
|
model string
|
|
mode string
|
|
timeout time.Duration
|
|
extraEnv []string
|
|
events chan core.Event
|
|
stdin io.WriteCloser
|
|
stdinMu sync.Mutex
|
|
closeOnce sync.Once
|
|
permReqID atomic.Value // stores string
|
|
chatID atomic.Value // stores string
|
|
ctx context.Context
|
|
cancel context.CancelFunc
|
|
wg sync.WaitGroup
|
|
alive atomic.Bool
|
|
}
|
|
|
|
var permissionPromptPattern = regexp.MustCompile(`(?is)(allow|approve|permission).{0,400}(\(y/n\)|\(y\/n\)|\(y\/N\)|\(Y\/n\)|\[y\/n\]|\[y\/N\]|\[Y\/n\]|yes\/no)`)
|
|
|
|
func newAntigravitySession(ctx context.Context, cmd, workDir, model, mode, resumeID string, extraEnv []string, timeout time.Duration) (*antigravitySession, error) {
|
|
sessionCtx, cancel := context.WithCancel(ctx)
|
|
|
|
as := &antigravitySession{
|
|
cmd: cmd,
|
|
workDir: workDir,
|
|
model: model,
|
|
mode: mode,
|
|
timeout: timeout,
|
|
extraEnv: extraEnv,
|
|
events: make(chan core.Event, 64),
|
|
ctx: sessionCtx,
|
|
cancel: cancel,
|
|
}
|
|
as.alive.Store(true)
|
|
|
|
if resumeID != "" && resumeID != core.ContinueSession {
|
|
as.chatID.Store(resumeID)
|
|
}
|
|
|
|
return as, nil
|
|
}
|
|
|
|
func (as *antigravitySession) Send(prompt string, images []core.ImageAttachment, files []core.FileAttachment) error {
|
|
if !as.alive.Load() {
|
|
return fmt.Errorf("session is closed")
|
|
}
|
|
|
|
// Capture existing chat logs so we can identify a new session on first turn
|
|
preEntries := make(map[string]bool)
|
|
homeDir, err := os.UserHomeDir()
|
|
if err == nil {
|
|
slug := antigravityProjectSlug(as.workDir)
|
|
chatsDir := filepath.Join(homeDir, ".gemini", "tmp", slug, "chats")
|
|
if entries, err := os.ReadDir(chatsDir); err == nil {
|
|
for _, entry := range entries {
|
|
if !entry.IsDir() && strings.HasSuffix(entry.Name(), ".jsonl") {
|
|
preEntries[entry.Name()] = true
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Save images and files into the workspace
|
|
attachDir := filepath.Join(as.workDir, ".cc-connect", "attachments")
|
|
if (len(images) > 0 || len(files) > 0) && os.MkdirAll(attachDir, 0o755) != nil {
|
|
attachDir = os.TempDir()
|
|
}
|
|
|
|
var imageRefs []string
|
|
for i, img := range images {
|
|
ext := ".png"
|
|
switch img.MimeType {
|
|
case "image/jpeg":
|
|
ext = ".jpg"
|
|
case "image/gif":
|
|
ext = ".gif"
|
|
case "image/webp":
|
|
ext = ".webp"
|
|
}
|
|
fname := fmt.Sprintf("img_%d_%d%s", time.Now().UnixMilli(), i, ext)
|
|
fpath := filepath.Join(attachDir, fname)
|
|
if err := os.WriteFile(fpath, img.Data, 0o600); err == nil {
|
|
imageRefs = append(imageRefs, fpath)
|
|
}
|
|
}
|
|
|
|
var fileRefs []string
|
|
for i, f := range files {
|
|
fname := filepath.Base(f.FileName)
|
|
if fname == "" || fname == "." || fname == ".." {
|
|
fname = fmt.Sprintf("file_%d_%d", time.Now().UnixMilli(), i)
|
|
}
|
|
fpath := filepath.Join(attachDir, fname)
|
|
if err := os.WriteFile(fpath, f.Data, 0o600); err == nil {
|
|
fileRefs = append(fileRefs, fpath)
|
|
}
|
|
}
|
|
|
|
chatID := as.CurrentSessionID()
|
|
isResume := chatID != ""
|
|
|
|
// Attach image and file references to prompt
|
|
fullPrompt := prompt
|
|
if len(imageRefs) > 0 {
|
|
if fullPrompt == "" {
|
|
fullPrompt = "Please analyze the attached image(s)."
|
|
}
|
|
fullPrompt += "\n\n[Attached images saved at: " + strings.Join(imageRefs, ", ") + "]"
|
|
}
|
|
if len(fileRefs) > 0 {
|
|
if fullPrompt == "" {
|
|
fullPrompt = "Please analyze the attached file(s)."
|
|
}
|
|
fullPrompt += "\n\n[Attached files saved at: " + strings.Join(fileRefs, ", ") + "]"
|
|
}
|
|
args := buildAntigravityArgs(chatID, isResume, as.mode, fullPrompt)
|
|
if strings.TrimSpace(as.model) != "" {
|
|
slog.Warn("antigravitySession: model is configured but ignored because agy does not support --model yet", "model", as.model)
|
|
}
|
|
|
|
var ctx context.Context
|
|
var cancel context.CancelFunc
|
|
if as.timeout > 0 {
|
|
ctx, cancel = context.WithTimeout(as.ctx, as.timeout)
|
|
} else {
|
|
ctx, cancel = context.WithCancel(as.ctx)
|
|
}
|
|
|
|
started := false
|
|
defer func() {
|
|
if !started {
|
|
cancel()
|
|
}
|
|
}()
|
|
|
|
slog.Debug("antigravitySession: launching", "resume", isResume, "args", core.RedactArgs(args))
|
|
cmd := exec.CommandContext(ctx, as.cmd, args...)
|
|
cmd.WaitDelay = 1 * time.Second
|
|
cmd.Dir = as.workDir
|
|
env := os.Environ()
|
|
if len(as.extraEnv) > 0 {
|
|
env = core.MergeEnv(env, as.extraEnv)
|
|
}
|
|
cmd.Env = env
|
|
|
|
stdout, err := cmd.StdoutPipe()
|
|
if err != nil {
|
|
return fmt.Errorf("antigravitySession: stdout pipe: %w", err)
|
|
}
|
|
var stdin io.WriteCloser
|
|
if usesInteractivePermission(as.mode) {
|
|
stdin, err = cmd.StdinPipe()
|
|
if err != nil {
|
|
return fmt.Errorf("antigravitySession: stdin pipe: %w", err)
|
|
}
|
|
}
|
|
|
|
var stderrBuf bytes.Buffer
|
|
cmd.Stderr = &stderrBuf
|
|
|
|
if err := cmd.Start(); err != nil {
|
|
return fmt.Errorf("antigravitySession: start: %w", err)
|
|
}
|
|
as.stdinMu.Lock()
|
|
as.stdin = stdin
|
|
as.stdinMu.Unlock()
|
|
|
|
started = true
|
|
as.wg.Add(1)
|
|
go func() {
|
|
defer cancel()
|
|
as.readLoop(ctx, cmd, stdout, &stderrBuf, append(imageRefs, fileRefs...), preEntries, time.Now())
|
|
}()
|
|
|
|
return nil
|
|
}
|
|
|
|
func buildAntigravityArgs(chatID string, isResume bool, mode, fullPrompt string) []string {
|
|
// Keep "-p <prompt>" at the very end because agy consumes the immediate next arg.
|
|
args := make([]string, 0, 10)
|
|
if isResume {
|
|
args = append(args, "--conversation", chatID)
|
|
}
|
|
switch mode {
|
|
case "yolo":
|
|
args = append(args, "--dangerously-skip-permissions")
|
|
case "plan":
|
|
args = append(args, "--sandbox")
|
|
}
|
|
args = append(args, "-p", fullPrompt)
|
|
return args
|
|
}
|
|
|
|
func usesInteractivePermission(mode string) bool {
|
|
return strings.EqualFold(strings.TrimSpace(mode), "default")
|
|
}
|
|
|
|
func (as *antigravitySession) readLoop(ctx context.Context, cmd *exec.Cmd, stdout io.ReadCloser, stderrBuf *bytes.Buffer, tempFiles []string, preEntries map[string]bool, sendStartedAt time.Time) {
|
|
defer as.wg.Done()
|
|
defer func() {
|
|
for _, f := range tempFiles {
|
|
_ = os.Remove(f)
|
|
}
|
|
|
|
// Detect conversation ID if this was the first turn of a fresh session
|
|
if as.CurrentSessionID() == "" {
|
|
var sid string
|
|
for attempt := 0; attempt < 15; attempt++ {
|
|
sid = as.detectNewSessionID(preEntries, sendStartedAt)
|
|
if sid != "" {
|
|
break
|
|
}
|
|
time.Sleep(20 * time.Millisecond)
|
|
}
|
|
if sid != "" {
|
|
as.chatID.Store(sid)
|
|
slog.Debug("antigravitySession: detected session ID", "session_id", sid)
|
|
// Emit an EventText carrying the session ID back to core
|
|
select {
|
|
case as.events <- core.Event{Type: core.EventText, SessionID: sid}:
|
|
case <-as.ctx.Done():
|
|
}
|
|
}
|
|
}
|
|
|
|
err := cmd.Wait()
|
|
sid := as.CurrentSessionID()
|
|
if err != nil {
|
|
stderrMsg := strings.TrimSpace(stderrBuf.String())
|
|
if stderrMsg != "" {
|
|
slog.Error("antigravitySession: process failed", "error", err, "stderr", stderrMsg)
|
|
select {
|
|
case as.events <- core.Event{Type: core.EventError, Error: fmt.Errorf("%s", stderrMsg)}:
|
|
case <-as.ctx.Done():
|
|
}
|
|
}
|
|
}
|
|
|
|
// Finalize turn
|
|
select {
|
|
case as.events <- core.Event{Type: core.EventResult, SessionID: sid, Done: true}:
|
|
case <-as.ctx.Done():
|
|
}
|
|
}()
|
|
|
|
go func() {
|
|
<-ctx.Done()
|
|
_ = stdout.Close()
|
|
}()
|
|
|
|
reader := bufio.NewReader(stdout)
|
|
buf := make([]byte, 1024)
|
|
permWindow := ""
|
|
|
|
for {
|
|
n, err := reader.Read(buf)
|
|
if n > 0 {
|
|
text := string(buf[:n])
|
|
permWindow += text
|
|
if len(permWindow) > 4096 {
|
|
permWindow = permWindow[len(permWindow)-4096:]
|
|
}
|
|
if pending, _ := as.permReqID.Load().(string); pending == "" {
|
|
if prompt, ok := extractPermissionPrompt(permWindow); ok {
|
|
requestID := fmt.Sprintf("agy-perm-%d", time.Now().UnixNano())
|
|
as.permReqID.Store(requestID)
|
|
select {
|
|
case as.events <- core.Event{
|
|
Type: core.EventPermissionRequest,
|
|
RequestID: requestID,
|
|
ToolName: "terminal_permission",
|
|
ToolInput: prompt,
|
|
ToolInputRaw: map[string]any{"prompt": prompt},
|
|
}:
|
|
case <-as.ctx.Done():
|
|
return
|
|
}
|
|
}
|
|
}
|
|
select {
|
|
case as.events <- core.Event{Type: core.EventText, Content: text}:
|
|
case <-as.ctx.Done():
|
|
return
|
|
}
|
|
}
|
|
if err != nil {
|
|
if err != io.EOF && !strings.Contains(err.Error(), "file already closed") {
|
|
slog.Error("antigravitySession: read error", "error", err)
|
|
select {
|
|
case as.events <- core.Event{Type: core.EventError, Error: fmt.Errorf("read stdout: %w", err)}:
|
|
case <-as.ctx.Done():
|
|
}
|
|
}
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (as *antigravitySession) detectNewSessionID(preEntries map[string]bool, sendStartedAt time.Time) string {
|
|
homeDir, err := os.UserHomeDir()
|
|
if err != nil {
|
|
return ""
|
|
}
|
|
slug := antigravityProjectSlug(as.workDir)
|
|
chatsDir := filepath.Join(homeDir, ".gemini", "tmp", slug, "chats")
|
|
|
|
entries, err := os.ReadDir(chatsDir)
|
|
if err != nil {
|
|
return ""
|
|
}
|
|
|
|
type candidate struct {
|
|
sessionID string
|
|
modTime time.Time
|
|
diff time.Duration
|
|
}
|
|
var candidates []candidate
|
|
for _, entry := range entries {
|
|
if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".jsonl") {
|
|
continue
|
|
}
|
|
if preEntries[entry.Name()] {
|
|
continue
|
|
}
|
|
info, err := entry.Info()
|
|
if err != nil {
|
|
continue
|
|
}
|
|
mod := info.ModTime()
|
|
if !sendStartedAt.IsZero() && mod.Before(sendStartedAt.Add(-2*time.Second)) {
|
|
continue
|
|
}
|
|
|
|
fpath := filepath.Join(chatsDir, entry.Name())
|
|
file, err := os.Open(fpath)
|
|
if err != nil {
|
|
continue
|
|
}
|
|
scanner := bufio.NewScanner(file)
|
|
scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
|
|
if scanner.Scan() {
|
|
var sf struct {
|
|
SessionID string `json:"sessionId"`
|
|
}
|
|
if json.Unmarshal([]byte(scanner.Text()), &sf) == nil && sf.SessionID != "" {
|
|
diff := time.Duration(0)
|
|
if !sendStartedAt.IsZero() {
|
|
if mod.After(sendStartedAt) {
|
|
diff = mod.Sub(sendStartedAt)
|
|
} else {
|
|
diff = sendStartedAt.Sub(mod)
|
|
}
|
|
}
|
|
candidates = append(candidates, candidate{
|
|
sessionID: sf.SessionID,
|
|
modTime: mod,
|
|
diff: diff,
|
|
})
|
|
}
|
|
}
|
|
_ = file.Close()
|
|
}
|
|
if len(candidates) == 0 {
|
|
return ""
|
|
}
|
|
sort.Slice(candidates, func(i, j int) bool {
|
|
if candidates[i].diff == candidates[j].diff {
|
|
return candidates[i].modTime.After(candidates[j].modTime)
|
|
}
|
|
return candidates[i].diff < candidates[j].diff
|
|
})
|
|
return candidates[0].sessionID
|
|
}
|
|
|
|
func extractPermissionPrompt(text string) (string, bool) {
|
|
loc := permissionPromptPattern.FindStringIndex(text)
|
|
if loc == nil {
|
|
return "", false
|
|
}
|
|
prompt := strings.TrimSpace(text[loc[0]:loc[1]])
|
|
if prompt == "" {
|
|
return "", false
|
|
}
|
|
return prompt, true
|
|
}
|
|
|
|
func (as *antigravitySession) RespondPermission(requestID string, result core.PermissionResult) error {
|
|
if !as.alive.Load() {
|
|
return fmt.Errorf("session is closed")
|
|
}
|
|
if pending, _ := as.permReqID.Load().(string); pending != "" && requestID != "" && requestID != pending {
|
|
return fmt.Errorf("permission request mismatch: got %q, pending %q", requestID, pending)
|
|
}
|
|
as.stdinMu.Lock()
|
|
defer as.stdinMu.Unlock()
|
|
if as.stdin == nil {
|
|
return fmt.Errorf("stdin is not available")
|
|
}
|
|
// agy permission prompts accept terminal-style responses.
|
|
// Keep this conservative until agy exposes a structured permission protocol.
|
|
reply := "y\n"
|
|
if strings.EqualFold(result.Behavior, "deny") {
|
|
reply = "n\n"
|
|
}
|
|
_, err := io.WriteString(as.stdin, reply)
|
|
if err != nil {
|
|
return fmt.Errorf("write permission response: %w", err)
|
|
}
|
|
as.permReqID.Store("")
|
|
return nil
|
|
}
|
|
|
|
func (as *antigravitySession) Events() <-chan core.Event {
|
|
return as.events
|
|
}
|
|
|
|
func (as *antigravitySession) CurrentSessionID() string {
|
|
v, _ := as.chatID.Load().(string)
|
|
return v
|
|
}
|
|
|
|
func (as *antigravitySession) Alive() bool {
|
|
return as.alive.Load()
|
|
}
|
|
|
|
func (as *antigravitySession) Close() error {
|
|
as.alive.Store(false)
|
|
as.cancel()
|
|
as.stdinMu.Lock()
|
|
if as.stdin != nil {
|
|
_ = as.stdin.Close()
|
|
as.stdin = nil
|
|
}
|
|
as.stdinMu.Unlock()
|
|
done := make(chan struct{})
|
|
go func() {
|
|
as.wg.Wait()
|
|
as.closeOnce.Do(func() {
|
|
close(as.events)
|
|
})
|
|
close(done)
|
|
}()
|
|
select {
|
|
case <-done:
|
|
case <-time.After(8 * time.Second):
|
|
slog.Warn("antigravitySession: close timed out")
|
|
}
|
|
return nil
|
|
}
|