iop/apps/node/internal/adapters/cli/status/antigravity.go

package status

import (
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"time"

	"github.com/creack/pty"
)

type AntigravityChecker struct {
	command string
}

func NewAntigravityChecker(command string) *AntigravityChecker {
	if command == "" {
		command = "agy"
	}
	return &AntigravityChecker{command: command}
}

func (g *AntigravityChecker) Check(ctx context.Context) (*UsageStatus, error) {
	cmd := exec.CommandContext(ctx, g.command, "--dangerously-skip-permissions")
	cmd.Env = append(os.Environ(), "TERM=xterm-256color")

	ptmx, err := pty.StartWithSize(cmd, &pty.Winsize{Rows: 40, Cols: 120})
	if err != nil {
		return nil, fmt.Errorf("failed to start pty: %w", err)
	}
	defer func() {
		_ = ptmx.Close()
		if cmd.Process != nil {
			_ = cmd.Process.Kill()
		}
	}()

	chunks := make(chan string, 100)
	go func() {
		buf := make([]byte, 4096)
		for {
			n, err := ptmx.Read(buf)
			if n > 0 {
				chunks <- string(buf[:n])
			}
			if err != nil {
				close(chunks)
				return
			}
		}
	}()

	var fullOutput string
	type waitPattern struct {
		name    string
		pattern *regexp.Regexp
	}

	waitForAny := func(startOffset int, patterns []waitPattern, timeout time.Duration) (string, error) {
		timeoutCh := time.After(timeout)
		for {
			targetOutput := ""
			if startOffset < len(fullOutput) {
				targetOutput = fullOutput[startOffset:]
			}
			cleanOutput := cleanANSI(targetOutput)
			for _, candidate := range patterns {
				if candidate.pattern.MatchString(cleanOutput) {
					return candidate.name, nil
				}
			}
			select {
			case <-ctx.Done():
				return "", ctx.Err()
			case <-timeoutCh:
				names := make([]string, 0, len(patterns))
				for _, candidate := range patterns {
					names = append(names, candidate.pattern.String())
				}
				return "", fmt.Errorf("timeout waiting for %s", strings.Join(names, " or "))
			case chunk, ok := <-chunks:
				if !ok {
					return "", io.EOF
				}
				fullOutput += chunk
			}
		}
	}

	sendText := func(text string) {
		for _, char := range text {
			_, _ = ptmx.Write([]byte{byte(char)})
			time.Sleep(50 * time.Millisecond)
		}
	}

	// 1. Wait for startup ready prompt.
	readyRegex := regexp.MustCompile(`for shortcuts|Type your message`)
	if _, err := waitForAny(0, []waitPattern{{name: "ready", pattern: readyRegex}}, 20*time.Second); err != nil {
		return nil, fmt.Errorf("antigravity startup failed: %w", err)
	}
	time.Sleep(300 * time.Millisecond)

	// 2. Request /usage for the overall quota.
	usageStart := len(fullOutput)
	sendText("\x15/usage\r")

	// 3. Wait for quota block, limit-reached, or no-api-calls notice.
	resultRegex := regexp.MustCompile(`(?i)Usage limit:|Limit reached|No API calls have been made|\d+(?:\.\d+)?%\s+used|\bquota\b`)
	if _, err := waitForAny(usageStart, []waitPattern{{name: "usage", pattern: resultRegex}}, 20*time.Second); err != nil {
		return nil, fmt.Errorf("failed waiting for antigravity /usage output: %w", err)
	}

	// 4. Request /model for per-model usage and reset times.
	time.Sleep(300 * time.Millisecond)
	modelStart := len(fullOutput)
	sendText("\x15/model\r")
	modelUsageRegex := regexp.MustCompile(`(?i)Model usage`)
	if _, err := waitForAny(modelStart, []waitPattern{{name: "model-usage", pattern: modelUsageRegex}}, 20*time.Second); err != nil {
		if ctx.Err() != nil {
			return nil, ctx.Err()
		}
		// Best-effort: ignore failure (timeout/EOF) and proceed to exit & parse what we have.
	}

	// Wait briefly for any trailing output.
	time.Sleep(500 * time.Millisecond)

	// Flush remaining chunks.
loop:
	for {
		select {
		case chunk, ok := <-chunks:
			if !ok {
				break loop
			}
			fullOutput += chunk
		default:
			break loop
		}
	}

	// 5. Graceful exit.
	sendText("\x1b") // Esc
	time.Sleep(100 * time.Millisecond)
	sendText("\x15") // Ctrl+U
	time.Sleep(100 * time.Millisecond)
	sendText("/quit\r")
	time.Sleep(300 * time.Millisecond)

	return ParseStatusOutput(fullOutput)
}