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

package status

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

	"github.com/creack/pty"
)

type CodexChecker struct {
	command string
}

func NewCodexChecker(command string) *CodexChecker {
	if command == "" {
		command = "codex"
	}
	return &CodexChecker{command: command}
}

func (c *CodexChecker) Check(ctx context.Context) (*UsageStatus, error) {
	cmd := exec.CommandContext(ctx, c.command, "--no-alt-screen")
	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(patterns []waitPattern, timeout time.Duration) (string, error) {
		timeoutCh := time.After(timeout)
		for {
			cleanOutput := cleanANSI(fullOutput)
			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
			}
		}
	}
	waitFor := func(pattern *regexp.Regexp, timeout time.Duration) error {
		_, err := waitForAny([]waitPattern{{name: pattern.String(), pattern: pattern}}, timeout)
		return err
	}

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

	// 1. Wait for startup ready prompt. Codex may show an update prompt before
	// the composer; choose "Skip" so the following slash command is not consumed.
	readyRegex := regexp.MustCompile(`(Tip:|model:)`)
	updatePromptRegex := regexp.MustCompile(`Update available!|Press enter to continue`)
	startupState, err := waitForAny([]waitPattern{
		{name: "ready", pattern: readyRegex},
		{name: "update", pattern: updatePromptRegex},
	}, 15*time.Second)
	if err != nil {
		return nil, fmt.Errorf("codex startup failed: %w", err)
	}
	if startupState == "update" {
		sendText("2\r")
		if err := waitFor(readyRegex, 15*time.Second); err != nil {
			return nil, fmt.Errorf("codex startup failed after skipping update prompt: %w", err)
		}
	}

	// 2. Request status
	sendText("\x15/status\r")

	// 3. Wait for either refresh message or actual limits
	statusOrRefreshRegex := regexp.MustCompile(`(?:refresh requested.*run /status again|Weekly limit:)`)
	if err := waitFor(statusOrRefreshRegex, 15*time.Second); err != nil {
		return nil, fmt.Errorf("failed waiting for first status: %w", err)
	}

	cleanOutput := cleanANSI(fullOutput)
	if strings.Contains(cleanOutput, "refresh requested") && !strings.Contains(cleanOutput, "Weekly limit:") {
		// Wait a moment for the backend to fetch new limits
		time.Sleep(2 * time.Second)

		// Clear current line (Ctrl+U) and re-request
		sendText("\x15/status\r")

		// Clear local buffer to ensure we find the *new* limits
		fullOutput = ""
		limitsRegex := regexp.MustCompile(`Weekly limit:`)
		if err := waitFor(limitsRegex, 15*time.Second); err != nil {
			return nil, fmt.Errorf("failed waiting for second status: %w", err)
		}
	}

	// Wait a tiny bit more to ensure progress bars are fully printed
	time.Sleep(500 * time.Millisecond)

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

	// Graceful exit
	sendText("/exit\r")
	time.Sleep(500 * time.Millisecond) // Give it time to exit cleanly

	return ParseStatusOutput(fullOutput)
}