147 lines
3.5 KiB
Go
147 lines
3.5 KiB
Go
package status
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"os/exec"
|
|
"regexp"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/creack/pty"
|
|
)
|
|
|
|
type GeminiChecker struct {
|
|
command string
|
|
}
|
|
|
|
func NewGeminiChecker(command string) *GeminiChecker {
|
|
if command == "" {
|
|
command = "gemini"
|
|
}
|
|
return &GeminiChecker{command: command}
|
|
}
|
|
|
|
func (g *GeminiChecker) Check(ctx context.Context) (*UsageStatus, error) {
|
|
cmd := exec.CommandContext(ctx, g.command, "--skip-trust", "--screen-reader")
|
|
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.
|
|
readyRegex := regexp.MustCompile(`for shortcuts|Type your message`)
|
|
if _, err := waitForAny([]waitPattern{{name: "ready", pattern: readyRegex}}, 20*time.Second); err != nil {
|
|
return nil, fmt.Errorf("gemini startup failed: %w", err)
|
|
}
|
|
time.Sleep(300 * time.Millisecond)
|
|
|
|
// 2. Request /stats model for the overall quota.
|
|
sendText("\x15/stats model\r")
|
|
|
|
// 3. Wait for quota block, limit-reached, or no-api-calls notice.
|
|
resultRegex := regexp.MustCompile(`(?i)Usage limit:|Limit reached|No API calls have been made`)
|
|
if _, err := waitForAny([]waitPattern{{name: "stats", pattern: resultRegex}}, 20*time.Second); err != nil {
|
|
return nil, fmt.Errorf("failed waiting for gemini /stats model output: %w", err)
|
|
}
|
|
|
|
// 4. Request /model for per-model usage and reset times.
|
|
time.Sleep(300 * time.Millisecond)
|
|
sendText("\x15/model\r")
|
|
modelUsageRegex := regexp.MustCompile(`(?i)Model usage`)
|
|
if _, err := waitForAny([]waitPattern{{name: "model-usage", pattern: modelUsageRegex}}, 20*time.Second); err != nil {
|
|
return nil, fmt.Errorf("failed waiting for gemini /model output: %w", err)
|
|
}
|
|
|
|
// Wait briefly for any trailing 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
|
|
}
|
|
}
|
|
|
|
// 5. Graceful exit.
|
|
sendText("\x1b") // Esc
|
|
time.Sleep(100 * time.Millisecond)
|
|
sendText("\x15") // Ctrl+U
|
|
time.Sleep(100 * time.Millisecond)
|
|
sendText("/quit\r")
|
|
time.Sleep(300 * time.Millisecond)
|
|
|
|
return ParseStatusOutput(fullOutput)
|
|
}
|