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) }