1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
// Varlink service exposing Claude plan usage.
//
// The protocol is hand-rolled (newline/NUL-framed JSON over a unix socket),
// mirroring the approach in ../maildir-varlink, so we avoid pulling in the
// varlink/go dependency.
//
// Interface: de.profpatsch.ClaudeUsage
//   method GetUsage() -> (
//     five_hour: ?Slot,
//     seven_day: ?Slot,
//     seven_day_sonnet: ?Slot,
//     seven_day_opus: ?Slot,
//     seven_day_cowork: ?Slot,
//     extra_usage_enabled: bool
//   )
//   type Slot (utilization: float, resets_at: string)
package main

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
)

const varlinkInterface = "de.profpatsch.ClaudeUsage"

// usageCache holds the most recent backend snapshot plus refresh metadata.
// GetUsage is served entirely from here; only the background refresher ever
// calls the backend.
type usageCache struct {
	mu          sync.RWMutex
	usage       *usageResponse // last *successful* fetch (may be nil if none yet)
	fetchedAt   time.Time      // time of last successful fetch
	reachable   bool           // whether the most recent attempt succeeded
	lastErr     error          // reason for the most recent failed attempt
	fingerprint string         // normalized signature of the last successful data
}

// refresher drives background fetches with exponential backoff.
type refresher struct {
	cache       *usageCache
	minInterval time.Duration
	maxInterval time.Duration
	curInterval time.Duration
	kick        chan struct{} // request an out-of-band fetch
}

var theRefresher *refresher

// varlinkCall is an incoming varlink method call.
type varlinkCall struct {
	Method     string         `json:"method"`
	Parameters map[string]any `json:"parameters,omitempty"`
	More       bool           `json:"more,omitempty"`
}

// varlinkReply is an outgoing varlink reply.
type varlinkReply struct {
	Parameters map[string]any `json:"parameters,omitempty"`
	Error      string         `json:"error,omitempty"`
}

const interfaceDescription = `# Claude plan usage limits, read from the claude.ai web API.
#
# The service maintains a cached snapshot refreshed by a background loop, so
# GetUsage never triggers a backend request on the hot path and always returns
# instantly. The refresh interval backs off exponentially while the data is
# unchanged, and resets to the minimum when the data changes. A GetUsage call
# also kicks an out-of-band refresh (unless one happened within the minimum
# interval) and resets the backoff.
interface de.profpatsch.ClaudeUsage

# A single usage window: a utilization percentage (0-100) and an ISO-8601
# timestamp at which the window resets.
type Slot (
  utilization: float,
  resets_at: string
)

# Returns the most recent cached usage windows. Any window the account does not
# have is returned as null.
#
# Metadata:
#   - fetched_at: ISO-8601 timestamp of the last *successful* backend fetch
#     (empty if none has succeeded yet).
#   - reachable: false if the most recent refresh attempt failed (network
#     error / unexpected status). When false, the usage fields still carry the
#     last good values (if any); clients should show a "stale/unreachable"
#     indicator. 'last_error' carries the reason.
#   - last_error: human-readable reason for the most recent refresh failure,
#     or empty when the last refresh succeeded.
method GetUsage() -> (
  five_hour: ?Slot,
  seven_day: ?Slot,
  seven_day_sonnet: ?Slot,
  seven_day_opus: ?Slot,
  seven_day_cowork: ?Slot,
  extra_usage_enabled: bool,
  fetched_at: string,
  reachable: bool,
  last_error: string
)

# The stored Firefox cookies were missing or rejected (stale sessionKey /
# cf_clearance), and no usable cached data exists yet. Reopen claude.ai in
# Firefox to refresh them. The 'hint' field describes the fix.
error StaleCredentials (message: string, hint: string)

# Any other failure fetching usage (network error, unexpected HTTP status,
# malformed response) when no cached data exists yet.
error FetchFailed (message: string, hint: string)
`

func runVarlinkServer(minInterval, maxInterval time.Duration) {
	socketPath := filepath.Join("/run/user", strconv.Itoa(os.Getuid()), "de.profpatsch.ClaudeUsage")

	// Remove a stale socket from a previous run.
	os.Remove(socketPath)

	listener, err := net.Listen("unix", socketPath)
	if err != nil {
		log.Fatalf("failed to listen on %s: %v\n", socketPath, err)
	}
	defer listener.Close()

	// Start the background refresher; GetUsage reads its cache.
	theRefresher = &refresher{
		cache:       &usageCache{},
		minInterval: minInterval,
		maxInterval: maxInterval,
		curInterval: minInterval,
		kick:        make(chan struct{}, 1),
	}
	go theRefresher.run()

	log.Printf("claude-usage varlink service listening on %s (refresh %s..%s)\n",
		socketPath, minInterval, maxInterval)

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Printf("accept error: %v\n", err)
			continue
		}
		go handleVarlinkConnection(conn)
	}
}

// run is the background refresh loop: fetch immediately, then wait either for
// the backoff timer or a kick, fetch again, and adjust the interval.
func (r *refresher) run() {
	r.fetch() // warm the cache on startup

	timer := time.NewTimer(r.curInterval)
	defer timer.Stop()

	for {
		select {
		case <-timer.C:
		case <-r.kick:
			if !timer.Stop() {
				<-timer.C
			}
		}
		r.fetch()
		timer.Reset(r.curInterval)
	}
}

// fetch performs one backend fetch and updates the cache + backoff interval.
//
//   - success, data unchanged: grow the interval (×2, capped at maxInterval).
//   - success, data changed:   reset the interval to minInterval.
//   - failure:                 keep last-good data, mark unreachable, and grow
//     the interval (so we don't hammer claude.ai on a persistent failure).
func (r *refresher) fetch() {
	usage, err := getUsage()

	r.cache.mu.Lock()
	defer r.cache.mu.Unlock()

	if err != nil {
		r.cache.reachable = false
		r.cache.lastErr = err
		r.curInterval = r.nextInterval()
		log.Printf("refresh failed (next in %s): %v\n", r.curInterval, err)
		return
	}

	fp := fingerprint(usage)
	changed := fp != r.cache.fingerprint

	r.cache.usage = usage
	r.cache.fetchedAt = time.Now()
	r.cache.reachable = true
	r.cache.lastErr = nil
	r.cache.fingerprint = fp

	if changed {
		r.curInterval = r.minInterval
	} else {
		r.curInterval = r.nextInterval()
	}
	log.Printf("refresh ok (changed=%v, next in %s)\n", changed, r.curInterval)
}

// nextInterval doubles the current interval, capped at maxInterval.
func (r *refresher) nextInterval() time.Duration {
	next := r.curInterval * 2
	if next > r.maxInterval {
		next = r.maxInterval
	}
	if next < r.minInterval {
		next = r.minInterval
	}
	return next
}

// requestKick triggers an out-of-band refresh unless the last successful fetch
// is newer than minInterval, and resets the backoff to the minimum so the
// service becomes responsive again after a client shows interest.
func (r *refresher) requestKick() {
	r.cache.mu.Lock()
	age := time.Since(r.cache.fetchedAt)
	r.cache.mu.Unlock()

	if age < r.minInterval {
		return
	}

	// Reset backoff so the kicked fetch (and subsequent cadence) is prompt.
	r.cache.mu.Lock()
	r.curInterval = r.minInterval
	r.cache.mu.Unlock()

	// Non-blocking: if a kick is already pending, that's fine.
	select {
	case r.kick <- struct{}{}:
	default:
	}
}

// fingerprint produces a normalized signature of the usage data that ignores
// the sub-minute jitter the backend adds to resets_at timestamps, so
// "unchanged" reflects an actual change in utilization or reset minute.
func fingerprint(u *usageResponse) string {
	var b strings.Builder
	add := func(name string, s *slot) {
		if s == nil {
			fmt.Fprintf(&b, "%s:nil|", name)
			return
		}
		fmt.Fprintf(&b, "%s:%.1f@%s|", name, s.Utilization, s.ResetsAt.UTC().Truncate(time.Minute).Format(time.RFC3339))
	}
	add("5h", u.FiveHour)
	add("7d", u.SevenDay)
	add("7ds", u.SevenDaySonnet)
	add("7do", u.SevenDayOpus)
	add("7dc", u.SevenDayCowork)
	fmt.Fprintf(&b, "extra:%v", u.ExtraUsage.IsEnabled)
	return b.String()
}

func handleVarlinkConnection(conn net.Conn) {
	defer conn.Close()

	scanner := bufio.NewScanner(conn)
	// Varlink frames messages with a trailing NUL byte.
	scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
		if i := strings.IndexByte(string(data), 0); i >= 0 {
			return i + 1, data[0:i], nil
		}
		if atEOF && len(data) > 0 {
			return len(data), data, nil
		}
		return 0, nil, nil
	})
	// Usage JSON can exceed the default 64KB token; be generous.
	scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)

	for scanner.Scan() {
		var req varlinkCall
		if err := json.Unmarshal(scanner.Bytes(), &req); err != nil {
			log.Printf("failed to parse request: %v\n", err)
			writeReply(conn, varlinkReply{Error: "org.varlink.service.InvalidParameter"})
			continue
		}

		reply := dispatchVarlink(req)
		if err := writeReply(conn, reply); err != nil {
			log.Printf("failed to write reply: %v\n", err)
			return
		}
	}
	if err := scanner.Err(); err != nil {
		log.Printf("scanner error: %v\n", err)
	}
}

func writeReply(conn net.Conn, reply varlinkReply) error {
	b, err := json.Marshal(reply)
	if err != nil {
		return err
	}
	b = append(b, 0) // NUL terminator
	_, err = conn.Write(b)
	return err
}

func dispatchVarlink(req varlinkCall) varlinkReply {
	switch req.Method {
	case "org.varlink.service.GetInfo":
		return varlinkReply{Parameters: map[string]any{
			"vendor":     "Profpatsch",
			"product":    "Claude Usage",
			"version":    "0.1",
			"url":        "none",
			"interfaces": []string{varlinkInterface},
		}}

	case "org.varlink.service.GetInterfaceDescription":
		iface, _ := req.Parameters["interface"].(string)
		if iface == varlinkInterface {
			return varlinkReply{Parameters: map[string]any{
				"description": interfaceDescription,
			}}
		}
		return varlinkReply{Error: "org.varlink.service.InterfaceNotFound"}

	case varlinkInterface + ".GetUsage":
		return handleGetUsage()

	default:
		log.Printf("unknown method: %s\n", req.Method)
		return varlinkReply{Error: "org.varlink.service.MethodNotFound"}
	}
}

func handleGetUsage() varlinkReply {
	// A client showing interest kicks an out-of-band refresh (unless one
	// happened recently) and resets the backoff. This returns the *current*
	// cached snapshot immediately; the kicked fetch lands on a later call.
	theRefresher.requestKick()

	c := theRefresher.cache
	c.mu.RLock()
	usage := c.usage
	fetchedAt := c.fetchedAt
	reachable := c.reachable
	lastErr := c.lastErr
	c.mu.RUnlock()

	// No usable cached data yet: surface the failure as an error, mirroring
	// the previous (uncached) behaviour.
	if usage == nil {
		err := lastErr
		if err == nil {
			err = fmt.Errorf("usage not yet available (service warming up)")
		}
		errName := varlinkInterface + ".FetchFailed"
		if errors.Is(err, errStaleCredentials) {
			errName = varlinkInterface + ".StaleCredentials"
		}
		return varlinkReply{
			Error: errName,
			Parameters: map[string]any{
				"message": err.Error(),
				"hint":    refreshHint,
			},
		}
	}

	slotParam := func(s *slot) any {
		if s == nil {
			return nil
		}
		return map[string]any{
			"utilization": s.Utilization,
			"resets_at":   s.ResetsAt.Format(timeRFC3339Nano),
		}
	}

	lastErrStr := ""
	if !reachable && lastErr != nil {
		lastErrStr = lastErr.Error()
	}
	fetchedAtStr := ""
	if !fetchedAt.IsZero() {
		fetchedAtStr = fetchedAt.UTC().Format(timeRFC3339Nano)
	}

	return varlinkReply{Parameters: map[string]any{
		"five_hour":           slotParam(usage.FiveHour),
		"seven_day":           slotParam(usage.SevenDay),
		"seven_day_sonnet":    slotParam(usage.SevenDaySonnet),
		"seven_day_opus":      slotParam(usage.SevenDayOpus),
		"seven_day_cowork":    slotParam(usage.SevenDayCowork),
		"extra_usage_enabled": usage.ExtraUsage.IsEnabled,
		"fetched_at":          fetchedAtStr,
		"reachable":           reachable,
		"last_error":          lastErrStr,
	}}
}

// timeRFC3339Nano is the layout used for resets_at in the varlink output.
const timeRFC3339Nano = "2006-01-02T15:04:05.999999999Z07:00"