133 lines
3 KiB
Go
133 lines
3 KiB
Go
package status
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"os/exec"
|
|
"regexp"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/creack/pty"
|
|
)
|
|
|
|
type CodexChecker struct{}
|
|
|
|
func NewCodexChecker() *CodexChecker {
|
|
return &CodexChecker{}
|
|
}
|
|
|
|
func (c *CodexChecker) Check(ctx context.Context) (*UsageStatus, error) {
|
|
cmd := exec.CommandContext(ctx, "codex", "--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
|
|
waitFor := func(pattern *regexp.Regexp, timeout time.Duration) error {
|
|
timeoutCh := time.After(timeout)
|
|
for {
|
|
if pattern.MatchString(cleanANSI(fullOutput)) {
|
|
return nil
|
|
}
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case <-timeoutCh:
|
|
return fmt.Errorf("timeout waiting for %s", pattern.String())
|
|
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) // fast typing
|
|
}
|
|
}
|
|
|
|
// 1. Wait for startup ready prompt
|
|
readyRegex := regexp.MustCompile(`(Tip:|model:)`)
|
|
if err := waitFor(readyRegex, 15*time.Second); err != nil {
|
|
return nil, fmt.Errorf("codex startup failed: %w", err)
|
|
}
|
|
|
|
// 2. Request status
|
|
sendText("/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(ctx, fullOutput)
|
|
}
|