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

package status

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

	"github.com/creack/pty"
)

type ClaudeChecker struct {
	command string
}

func NewClaudeChecker(command string) *ClaudeChecker {
	if command == "" {
		command = "claude"
	}
	return &ClaudeChecker{command: command}
}

func tailLines(s string, n int) string {
	lines := strings.Split(s, "\n")
	if len(lines) > n {
		lines = lines[len(lines)-n:]
	}
	return strings.Join(lines, "\n")
}

func (c *ClaudeChecker) Check(ctx context.Context) (*UsageStatus, error) {
	cmd := exec.CommandContext(ctx, c.command)
	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
			}
		}
	}
	sendText := func(text string) {
		for _, char := range text {
			_, _ = ptmx.Write([]byte{byte(char)})
			time.Sleep(50 * time.Millisecond)
		}
	}

	// 1. Wait for startup ready prompt. The shortcut hint bar ("? for shortcuts"
	// or "/effort") only renders after the TUI is fully drawn and ready for
	// input — matching that prevents sending /usage before keystrokes are accepted.
	readyRegex := regexp.MustCompile(`for shortcuts|/effort|Welcome back|Try /`)
	if _, err := waitForAny([]waitPattern{{name: "ready", pattern: readyRegex}}, 20*time.Second); err != nil {
		return nil, fmt.Errorf("claude startup failed: %w", err)
	}
	time.Sleep(500 * time.Millisecond)

	// 2. Request Claude's status panel. This is intentionally the same command
	// path for both the headless "claude" profile and the persistent
	// "claude-tui" profile; the checker starts a short-lived TUI solely to read
	// status output.
	sendText("\x15/status\r")

	// 3. Wait until the status screen is visible, then wait for the terminal
	// output to settle. Older Claude builds/tests may still render usage-style
	// "Current session / Current week" blocks, so those markers remain accepted.
	statusPanelRegex := regexp.MustCompile(`(?i)Version\s*:|Session\s*ID\s*:|Login\s+method\s*:|Current\s+session|Current\s+week|Esc\s+to\s+cancel|d\s+to\s+day`)
	if _, err := waitForAny([]waitPattern{{name: "status-panel", pattern: statusPanelRegex}}, 20*time.Second); err != nil {
		return nil, fmt.Errorf("claude /status output did not open: %w", err)
	}

	waitForQuietOutput := func(idleTimeout, maxWait time.Duration) error {
		idleCh := time.NewTimer(idleTimeout)
		defer idleCh.Stop()
		maxCh := time.NewTimer(maxWait)
		defer maxCh.Stop()
		for {
			if st, _ := ParseStatusOutput(fullOutput); st != nil && st.DailyLimit != "" && st.WeeklyLimit != "" {
				return nil
			}
			select {
			case <-ctx.Done():
				return ctx.Err()
			case <-maxCh.C:
				return nil
			case <-idleCh.C:
				return nil
			case chunk, ok := <-chunks:
				if !ok {
					return io.EOF
				}
				fullOutput += chunk
				if !idleCh.Stop() {
					select {
					case <-idleCh.C:
					default:
					}
				}
				idleCh.Reset(idleTimeout)
			}
		}
	}
	if err := waitForQuietOutput(1500*time.Millisecond, 10*time.Second); err != nil {
		if st, parseErr := ParseStatusOutput(fullOutput); parseErr == nil && st != nil && strings.TrimSpace(st.RawOutput) != "" {
			return st, nil
		}
		return nil, fmt.Errorf("failed waiting for status output: %w", err)
	}

	// Wait a bit more for full 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
		}
	}

	// 4. Graceful exit attempt
	sendText("\x1b") // Esc
	time.Sleep(100 * time.Millisecond)
	sendText("\x15") // Ctrl+U (clear line)
	time.Sleep(100 * time.Millisecond)
	sendText("/exit\r")
	time.Sleep(500 * time.Millisecond)

	return ParseStatusOutput(fullOutput)
}