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 /usage
	sendText("\x15/usage\r")

	// 3a. Wait for the /usage panel itself to open (the panel tab strip contains
	// "Usage" highlighted between "Status/Config/.../Stats"). This separates
	// "command palette didn't accept /usage" from "panel opened but data hasn't
	// rendered" — they have different error tails.
	usagePanelRegex := regexp.MustCompile(`(?i)Current\s+session|Current\s+week|Esc\s+to\s+cancel|d\s+to\s+day`)
	if _, err := waitForAny([]waitPattern{{name: "usage-panel", pattern: usagePanelRegex}}, 15*time.Second); err != nil {
		return nil, fmt.Errorf("claude /usage panel did not open: %w", err)
	}

	// 3b. Wait for both session and week data to parse from the terminal
	// visible screen. We do not send `d`/`w` nudges: in real Claude they
	// transition the lower panel between Last-24h and Last-7d views and can
	// shift the Current session / Current week block off the visible region
	// mid-render, which has been observed to corrupt the screen state right
	// when the parser would have otherwise succeeded. The screen-aware
	// parser handles cursor repaint without any input from us.
	fullUsageWait := func(timeout time.Duration) error {
		timeoutCh := time.After(timeout)
		for {
			if st, _ := ParseStatusOutput(fullOutput); st != nil && st.DailyLimit != "" && st.WeeklyLimit != "" {
				return nil
			}
			select {
			case <-ctx.Done():
				return ctx.Err()
			case <-timeoutCh:
				visTail := tailLines(renderVisibleScreen(fullOutput, 60, 200), 40)
				rawTail := tailLines(cleanANSI(fullOutput), 40)
				return fmt.Errorf("timeout waiting for complete Claude usage block\nvisible-screen tail:\n%s\nraw-clean tail:\n%s", visTail, rawTail)
			case chunk, ok := <-chunks:
				if !ok {
					return io.EOF
				}
				fullOutput += chunk
			}
		}
	}
	if err := fullUsageWait(60 * time.Second); err != nil {
		return nil, fmt.Errorf("failed waiting for usage 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)
}