interactive_codex_usage.go RawHistoryFolder

1package modes
2
3import (
4	"context"
5	"fmt"
6	"time"
7
8	"github.com/bnema/zut/packages/core"
9	"github.com/bnema/zut/packages/provider"
10)
11
12const codexUsageInterval = 5 * time.Minute
13
14// State is protected by i.mu. Each request has
15// its own buffered result channel, so canceled/account-switched requests can
16// finish without blocking or publishing into the new account's display.
17type codexUsageState struct {
18	next   time.Time
19	usage  *provider.CodexWeeklyUsage
20	result <-chan *provider.CodexWeeklyUsage
21	cancel context.CancelFunc
22}
23
24type codexUsageClient struct {
25	provider.Client
26	fetch func(context.Context) (*provider.CodexWeeklyUsage, error)
27}
28
29// BindCodexWeeklyUsageFetcher attaches the optional meter to a private candidate
30// before publication. Preparing or discarding it cannot affect the live meter.
31// The host supplies nil for API keys, custom endpoints, and other providers.
32func BindCodexWeeklyUsageFetcher(ag *core.Agent, fetch func(context.Context) (*provider.CodexWeeklyUsage, error)) {
33	if fetch != nil {
34		ag.Client = &codexUsageClient{Client: ag.Client, fetch: fetch}
35	}
36}
37
38func (i *Interactive) codexUsageFetcherLocked() func(context.Context) (*provider.CodexWeeklyUsage, error) {
39	if i.agent != nil && i.cfg.Provider == "openai-codex" {
40		if client, ok := i.agent.Client.(*codexUsageClient); ok {
41			return client.fetch
42		}
43	}
44	return nil
45}
46
47func (i *Interactive) resetCodexUsage() {
48	i.mu.Lock()
49	defer i.mu.Unlock()
50	i.resetCodexUsageLocked()
51}
52
53func (i *Interactive) resetCodexUsageLocked() {
54	if i.codexUsage.cancel != nil {
55		i.codexUsage.cancel()
56	}
57	i.codexUsage = codexUsageState{}
58}
59
60func (i *Interactive) refreshCodexUsage(ctx context.Context, now time.Time) {
61	i.mu.Lock()
62	defer i.mu.Unlock()
63	fetch := i.codexUsageFetcherLocked()
64	if fetch == nil || ctx.Err() != nil {
65		i.resetCodexUsageLocked()
66		return
67	}
68	state := &i.codexUsage
69	select {
70	case usage := <-state.result:
71		state.usage = usage
72		state.result = nil
73		state.cancel()
74		state.cancel = nil
75		i.invalidate()
76	default:
77	}
78	if state.result != nil || now.Before(state.next) {
79		return
80	}
81	state.next = now.Add(codexUsageInterval)
82	state.usage = nil
83	i.invalidate()
84	requestCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
85	state.cancel = cancel
86	result := make(chan *provider.CodexWeeklyUsage, 1)
87	state.result = result
88	go func() {
89		defer cancel()
90		usage, err := fetch(requestCtx)
91		if err != nil {
92			// Background usage is best-effort; never interrupt a conversation
93			// or expose credentials/account payloads in the status error area.
94			usage = nil
95		}
96		result <- usage
97	}()
98}
99
100func (i *Interactive) codexWeeklyLabelLocked(now time.Time) string {
101	if i.codexUsageFetcherLocked() == nil {
102		return ""
103	}
104	usage := i.codexUsage.usage
105	if usage == nil ||
106		!now.Before(i.codexUsage.next) ||
107		(!usage.ResetsAt.IsZero() && !now.Before(usage.ResetsAt)) {
108		return "weekly:?"
109	}
110	return fmt.Sprintf("weekly:%.0f%%", usage.RemainingPercent)
111}