From f8e5a68a49bac13f0cb27fd6b8a67547cd0e4ad3 Mon Sep 17 00:00:00 2001 From: StuBehan Date: Tue, 22 Sep 2026 08:24:53 +0100 Subject: [PATCH 1/3] feat(usage): track a fair-usage budget for pi on the Usage tab --- README.md | 4 +- .../PiUsageTests.swift | 287 ++++++++++++++++++ .../WidgetQuotaTests.swift | 40 ++- notify.conf.example | 10 + panel/Panel.swift | 13 + panel/PanelNav.swift | 46 ++- panel/PiUsage.swift | 133 ++++++++ panel/SessionUsage.swift | 32 +- panel/Settings.swift | 4 + panel/UsageHistory.swift | 116 ++++++- panel/WidgetQuota.swift | 16 +- 11 files changed, 678 insertions(+), 23 deletions(-) create mode 100644 Tests/StackNudgePanelCoreTests/PiUsageTests.swift create mode 100644 panel/PiUsage.swift diff --git a/README.md b/README.md index 3525752..7840b2f 100644 --- a/README.md +++ b/README.md @@ -257,10 +257,12 @@ Codex works the same way through a different door: stack-nudge asks the local `c Polls every 60 seconds while the full panel is open, and otherwise every 5 minutes by default (configurable via Settings → Usage → "Poll frequency"). The collapsed widget counts as background, so it polls at your configured frequency rather than the faster open-panel rate. Opening the panel syncs immediately if the last one is over a minute old. On the Usage tab: `r` triggers a manual sync, `p` pauses/resumes the poller. -Claude, Codex, and Antigravity each appear in the tab's client list when they have quota to show; `↑`/`↓` switch between them. **The compact widget's gauge follows whatever you select here** — pick Codex in the Usage tab and the pill's rings, hover legend, and reset countdown all switch to Codex, with the client name shown in the legend on hover. All three are read on the same poll tick, so switching costs nothing. The selection is in-memory and resets to the first connected client on relaunch. +Claude, Codex, Antigravity and pi each appear in the tab's client list when they have something to show; `↑`/`↓` switch between them. **The compact widget's gauge follows whatever you select here** — pick Codex in the Usage tab and the pill's rings, hover legend, and reset countdown all switch to Codex, with the client name shown in the legend on hover. All four are read on the same poll tick, so switching costs nothing. The selection is in-memory and resets to the first connected client on relaunch. For Claude and Codex the two rings are the 5-hour and weekly windows — though Codex reports the window length per limit rather than a fixed pair, and on some accounts publishes only a weekly one, so its Usage tab headings are named from what it actually reports. Antigravity reports neither — it publishes one window per model plus a monthly credit pool — so its inner ring shows whichever model is closest to its limit and the outer ring shows monthly prompt credits. Gemini CLI has no usage counter: unlike the others it writes no rate-limit data to disk and serves no local endpoint, so there's nothing to read. +**pi is the one client whose bars are not a quota.** It enforces none: its API models bill per token against your own keys, and its local models cost nothing. What its rows track is a fair-usage budget you set yourself — a daily token allowance per lane, with the week being seven of them — so every row names its denominator (*"of 1.5M budget"*) and the tab labels it `Budget` where the others show a subscription tier. The windows are calendar ones, today and this week, which is what makes the reset time and the ahead-of-pace warning mean anything. Tokens are counted the way the history graph's "no cache reads" metric counts them, so resuming a long session doesn't burn the budget by replaying itself, and a turn counts as local when pi priced it at zero. Set the two allowances in Settings → Usage (`STACKNUDGE_PI_API_BUDGET` / `STACKNUDGE_PI_LOCAL_BUDGET`, tokens per day, `Off` to drop a lane's rows). Going over is reported rather than clamped: that is the only thing a budget you set for yourself can usefully tell you. + #### Threshold-crossing notifications When any quota tier reaches your configured threshold, stack-nudge fires a banner — *"Weekly quota at 85% — resets May 17"* — once per period per tier, so you get a heads-up before hitting the cap. Configure in Settings → Usage: diff --git a/Tests/StackNudgePanelCoreTests/PiUsageTests.swift b/Tests/StackNudgePanelCoreTests/PiUsageTests.swift new file mode 100644 index 0000000..ff734d2 --- /dev/null +++ b/Tests/StackNudgePanelCoreTests/PiUsageTests.swift @@ -0,0 +1,287 @@ +import XCTest + +@testable import StackNudgePanelCore + +// Pi's row on the Usage tab is a budget the user sets, not a quota a provider +// enforces, so what these pin is that every number under it is derived rather +// than assumed: which turns count as local, how wide the windows really are, +// and that a week's worth of history survives the graph's narrower refresh. +final class PiUsageTests: XCTestCase { + + // MARK: Fixtures + + private func fixtureDirectory() -> String { + let dir = NSTemporaryDirectory() + "pi-usage-\(UUID().uuidString)/" + try? FileManager.default.createDirectory(atPath: dir, withIntermediateDirectories: true) + return dir + } + + @discardableResult + private func write(_ lines: [String], to directory: String, name: String, modified: Date) -> String { + let path = directory + name + try? (lines.joined(separator: "\n") + "\n").write(toFile: path, atomically: true, encoding: .utf8) + try? FileManager.default.setAttributes([.modificationDate: modified], ofItemAtPath: path) + return path + } + + private static let stamp: ISO8601DateFormatter = { + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + return formatter + }() + + // A real-shape pi assistant message. `cost.total` is the only thing that + // separates a billed API model from a free local one — pi writes the same + // usage block for both. + private func piLine(at when: Date, + input: Int = 1_000, + output: Int = 100, + cacheRead: Int = 5_000, + cacheWrite: Int = 0, + reasoning: Int = 0, + cost: Double = 0, + id: String = UUID().uuidString) -> String { + """ + {"type":"message","id":"\(id)","timestamp":"\(Self.stamp.string(from: when))",\ + "message":{"role":"assistant","model":"m","usage":{"input":\(input),\ + "output":\(output),"cacheRead":\(cacheRead),"cacheWrite":\(cacheWrite),\ + "reasoning":\(reasoning),"totalTokens":0,"cost":{"total":\(cost)}}}} + """ + } + + private func totals(_ lines: [String], now: Date, span: TimeInterval = 8 * 86400) -> UsageTotals { + let dir = fixtureDirectory() + write(lines, to: dir, name: "session.jsonl", modified: now) + let store = UsageHistoryStore() + store.refresh(source: .pi, root: dir, now: now, retaining: span) + return store.totals(source: .pi, from: now.addingTimeInterval(-span), to: now.addingTimeInterval(60)) + } + + // MARK: Local vs API + + func test_zeroCostTurnCountsAsLocal() { + let now = Date(timeIntervalSince1970: 1_785_150_000) + let actual = totals([piLine(at: now.addingTimeInterval(-60), cost: 0)], now: now) + XCTAssertEqual(actual.localTokens, 1_100) + XCTAssertEqual(actual.apiTokens, 0) + } + + // pi writes a local turn's cost as the integer 0, a billed one as a + // fraction. A cast that only accepts one of those shapes silently files + // every turn in the wrong lane. + func test_integerCostIsReadAsANumber() { + let now = Date(timeIntervalSince1970: 1_785_150_000) + let dir = fixtureDirectory() + let line = """ + {"type":"message","id":"a","timestamp":"\(Self.stamp.string(from: now.addingTimeInterval(-60)))",\ + "message":{"role":"assistant","usage":{"input":10,"output":20,"cacheRead":0,\ + "cacheWrite":0,"reasoning":0,"cost":{"total":2}}}} + """ + write([line], to: dir, name: "session.jsonl", modified: now) + let store = UsageHistoryStore() + store.refresh(source: .pi, root: dir, now: now, retaining: 86400) + + let actual = store.totals(source: .pi, from: now.addingTimeInterval(-86400), to: now) + XCTAssertEqual(actual.apiTokens, 30) + XCTAssertEqual(actual.localTokens, 0) + } + + func test_pricedTurnCountsAsApi() { + let now = Date(timeIntervalSince1970: 1_785_150_000) + let actual = totals([piLine(at: now.addingTimeInterval(-60), cost: 0.0112)], now: now) + XCTAssertEqual(actual.apiTokens, 1_100) + XCTAssertEqual(actual.localTokens, 0) + } + + // Reasoning is output the model produced; cache reads are a replay of + // content already paid for and must not burn a budget. + func test_reasoningCountsAndCacheReadsDoNot() { + let now = Date(timeIntervalSince1970: 1_785_150_000) + let actual = totals([piLine(at: now.addingTimeInterval(-60), + input: 10, output: 20, cacheRead: 9_000, + cacheWrite: 5, reasoning: 30, cost: 1)], + now: now) + XCTAssertEqual(actual.apiTokens, 65) + } + + func test_bothLanesInOneWindow() { + let now = Date(timeIntervalSince1970: 1_785_150_000) + let actual = totals([piLine(at: now.addingTimeInterval(-60), cost: 1), + piLine(at: now.addingTimeInterval(-120), cost: 0)], + now: now) + XCTAssertEqual(actual.apiTokens, 1_100) + XCTAssertEqual(actual.localTokens, 1_100) + XCTAssertEqual(actual.turns, 2) + } + + // A resumed session replays earlier turns into a second transcript. + func test_replayedTurnCountsOnce() { + let now = Date(timeIntervalSince1970: 1_785_150_000) + let line = piLine(at: now.addingTimeInterval(-60), cost: 1, id: "repeated") + let actual = totals([line, line], now: now) + XCTAssertEqual(actual.turns, 1) + XCTAssertEqual(actual.apiTokens, 1_100) + } + + func test_nonAssistantLinesIgnored() { + let now = Date(timeIntervalSince1970: 1_785_150_000) + let session = #"{"type":"session","version":3,"id":"s","cwd":"/tmp"}"# + let user = """ + {"type":"message","id":"u","timestamp":"\(Self.stamp.string(from: now))",\ + "message":{"role":"user","content":[{"type":"text","text":"ask the assistant"}]}} + """ + let actual = totals([session, user, piLine(at: now.addingTimeInterval(-60), cost: 1)], now: now) + XCTAssertEqual(actual.turns, 1) + } + + // MARK: Window arithmetic + + // The span a totals() caller asks for is honoured exactly, rather than being + // silently truncated to the graph's widest bucket grid. + func test_totalsSpanAFullWeek() { + let now = Date(timeIntervalSince1970: 1_785_150_000) + let actual = totals([piLine(at: now.addingTimeInterval(-6 * 86400), cost: 1), + piLine(at: now.addingTimeInterval(-60), cost: 1)], + now: now) + XCTAssertEqual(actual.turns, 2) + XCTAssertEqual(actual.apiTokens, 2_200) + } + + func test_entriesOutsideTheSpanAreExcluded() { + let now = Date(timeIntervalSince1970: 1_785_150_000) + let dir = fixtureDirectory() + write([piLine(at: now.addingTimeInterval(-3 * 86400), cost: 1), + piLine(at: now.addingTimeInterval(-60), cost: 1)], + to: dir, name: "session.jsonl", modified: now) + let store = UsageHistoryStore() + store.refresh(source: .pi, root: dir, now: now, retaining: 8 * 86400) + + let actual = store.totals(source: .pi, from: now.addingTimeInterval(-86400), to: now) + XCTAssertEqual(actual.turns, 1) + } + + // The graph refreshes the same source over 24h. Without a retention + // watermark that narrower call evicts the budget's week between polls. + func test_narrowRefreshDoesNotEvictTheWeek() { + let now = Date(timeIntervalSince1970: 1_785_150_000) + let dir = fixtureDirectory() + let old = now.addingTimeInterval(-5 * 86400) + write([piLine(at: old, cost: 1)], to: dir, name: "old.jsonl", modified: old) + let store = UsageHistoryStore() + store.refresh(source: .pi, root: dir, now: now, retaining: 8 * 86400) + store.refresh(source: .pi, root: dir, now: now, retaining: UsageWindow.widest.seconds) + + let actual = store.totals(source: .pi, from: now.addingTimeInterval(-7 * 86400), to: now) + XCTAssertEqual(actual.apiTokens, 1_100) + } + + // MARK: Budget + + private func window(_ start: TimeInterval, _ duration: TimeInterval) -> DateInterval { + DateInterval(start: Date(timeIntervalSince1970: start), duration: duration) + } + + func test_utilizationIsTokensOverBudget() { + let day = window(1_785_110_400, 86400) + let snapshot = PiUsageBudget.snapshot(day: day, + dayTotals: UsageTotals(localTokens: 0, apiTokens: 750_000, turns: 1), + week: day, + weekTotals: UsageTotals(localTokens: 0, apiTokens: 750_000, turns: 1), + budget: PiBudget(apiDaily: 1_500_000, localDaily: 500_000)) + XCTAssertEqual(snapshot?.apiToday?.utilization, 50) + XCTAssertEqual(snapshot?.apiToday?.resetsAt, day.end) + XCTAssertEqual(snapshot?.apiToday?.windowLength, 86400) + XCTAssertNil(snapshot?.localToday) + } + + // Going over is the thing a self-imposed budget exists to report. + func test_overBudgetIsNotClamped() { + let day = window(1_785_110_400, 86400) + let snapshot = PiUsageBudget.snapshot(day: day, + dayTotals: UsageTotals(localTokens: 3_000_000, apiTokens: 0, turns: 1), + week: day, + weekTotals: UsageTotals(localTokens: 3_000_000, apiTokens: 0, turns: 1), + budget: PiBudget(apiDaily: 1_500_000, localDaily: 500_000)) + XCTAssertEqual(snapshot?.localToday?.utilization, 600) + } + + func test_weeklyBudgetIsSevenDays() { + XCTAssertEqual(PiBudget(apiDaily: 1_000_000, localDaily: 100_000).apiWeekly, 7_000_000) + XCTAssertEqual(PiBudget(apiDaily: 1_000_000, localDaily: 100_000).localWeekly, 700_000) + } + + // A lane switched off in Settings draws no row, however much it was used. + func test_zeroBudgetDrawsNoTier() { + let day = window(1_785_110_400, 86400) + let snapshot = PiUsageBudget.snapshot(day: day, + dayTotals: UsageTotals(localTokens: 400_000, apiTokens: 0, turns: 1), + week: day, + weekTotals: UsageTotals(localTokens: 400_000, apiTokens: 0, turns: 1), + budget: PiBudget(apiDaily: 1_500_000, localDaily: 0)) + XCTAssertNil(snapshot) + } + + func test_noUsageProducesNoSnapshot() { + let day = window(1_785_110_400, 86400) + XCTAssertNil(PiUsageBudget.snapshot(day: day, dayTotals: UsageTotals(), + week: day, weekTotals: UsageTotals(), + budget: .fallback)) + } + + // MARK: Probe + + // End to end: fixtures on disk through the probe to a snapshot, with the + // day boundary separating what counts today from what only counts this week. + func test_probeSplitsTodayFromTheWeek() throws { + var calendar = Calendar(identifier: .gregorian) + calendar.timeZone = TimeZone(identifier: "UTC")! + let now = calendar.date(from: DateComponents(year: 2026, month: 9, day: 23, hour: 9))! + let dir = fixtureDirectory() + write([piLine(at: now.addingTimeInterval(-3_600), input: 990, output: 10, cost: 1), + piLine(at: now.addingTimeInterval(-24 * 3_600), input: 500, output: 0, cost: 1)], + to: dir, name: "session.jsonl", modified: now) + + let actual = PiUsageProbe.read(store: UsageHistoryStore(), + root: dir, + budget: PiBudget(apiDaily: 1_000_000, localDaily: 500_000), + now: now, + calendar: calendar) + + XCTAssertEqual(actual?.apiToday?.utilization, 0.1) + let week = try XCTUnwrap(actual?.apiThisWeek?.utilization) + XCTAssertEqual(week, 1_500.0 / 7_000_000 * 100, accuracy: 0.000_001) + XCTAssertNil(actual?.localToday) + } + + // MARK: Calendar windows + + func test_windowsCoverNow() { + let now = Date(timeIntervalSince1970: 1_785_150_000) + let windows = PiUsageBudget.windows(now: now) + XCTAssertEqual(windows?.day.contains(now), true) + XCTAssertEqual(windows?.week.contains(now), true) + XCTAssertEqual(windows?.day.duration, 86400) + } + + // A spring-forward day is 23 hours long. The pane's ahead-of-pace warning + // divides by this, so measuring the period beats assuming 86400. + func test_dstDayIsShorterThanTwentyFourHours() { + var calendar = Calendar(identifier: .gregorian) + calendar.timeZone = TimeZone(identifier: "Europe/London")! + let noon = calendar.date(from: DateComponents(year: 2026, month: 3, day: 29, hour: 12))! + + let windows = PiUsageBudget.windows(now: noon, calendar: calendar) + XCTAssertEqual(windows?.day.duration, 23 * 3600) + } + + func test_weekStartsOnTheCalendarsFirstWeekday() { + var calendar = Calendar(identifier: .gregorian) + calendar.timeZone = TimeZone(identifier: "UTC")! + calendar.firstWeekday = 2 // Monday + let wednesday = calendar.date(from: DateComponents(year: 2026, month: 9, day: 23, hour: 9))! + + let windows = PiUsageBudget.windows(now: wednesday, calendar: calendar) + XCTAssertEqual(calendar.component(.weekday, from: windows!.week.start), 2) + XCTAssertEqual(windows?.week.duration, 7 * 86400) + } +} diff --git a/Tests/StackNudgePanelCoreTests/WidgetQuotaTests.swift b/Tests/StackNudgePanelCoreTests/WidgetQuotaTests.swift index c44d53b..9ad4761 100644 --- a/Tests/StackNudgePanelCoreTests/WidgetQuotaTests.swift +++ b/Tests/StackNudgePanelCoreTests/WidgetQuotaTests.swift @@ -36,12 +36,48 @@ final class WidgetQuotaTests: XCTestCase { private func make(_ client: UsageClient?, claude: QuotaSnapshot? = nil, codex: CodexQuotaSnapshot? = nil, - agy: AntigravityQuotaSnapshot? = nil) -> WidgetQuota { - WidgetQuota.make(client: client, claude: claude, codex: codex, antigravity: agy) + agy: AntigravityQuotaSnapshot? = nil, + pi: PiQuotaSnapshot? = nil) -> WidgetQuota { + WidgetQuota.make(client: client, claude: claude, codex: codex, antigravity: agy, pi: pi) + } + + private func piSnapshot(apiToday: Double?, apiWeek: Double?, + localToday: Double?, localWeek: Double?) -> PiQuotaSnapshot { + let day = DateInterval(start: Date(), duration: 86400) + let week = DateInterval(start: Date(), duration: 7 * 86400) + func tier(_ used: Double?, _ window: DateInterval) -> QuotaTier? { + used.map { QuotaTier(utilization: $0, resetsAt: window.end, windowLength: window.duration) } + } + return PiQuotaSnapshot(apiToday: tier(apiToday, day), + apiThisWeek: tier(apiWeek, week), + localToday: tier(localToday, day), + localThisWeek: tier(localWeek, week), + budget: .fallback) } // MARK: - Per-client ring mapping + func test_pi_mapsTodayAndThisWeek() { + let q = make(.pi, pi: piSnapshot(apiToday: 62, apiWeek: 18, localToday: 4, localWeek: 2)) + XCTAssertEqual(q.short?.utilization, 62) + XCTAssertEqual(q.long?.utilization, 18) + XCTAssertEqual(q.shortLabel, "1d") + XCTAssertEqual(q.longLabel, "7d") + } + + // Local-only usage still gets both rings rather than falling back to empty. + func test_pi_fallsBackToLocalWhenNoApiUsage() { + let q = make(.pi, pi: piSnapshot(apiToday: nil, apiWeek: nil, localToday: 30, localWeek: 9)) + XCTAssertEqual(q.short?.utilization, 30) + XCTAssertEqual(q.long?.utilization, 9) + } + + // A budget is the user's own, so passing it is the point of the row. + func test_pi_overBudgetIsNotClamped() { + let q = make(.pi, pi: piSnapshot(apiToday: 140, apiWeek: 30, localToday: nil, localWeek: nil)) + XCTAssertEqual(q.short?.utilization, 140) + } + func test_claude_mapsFiveHourAndSevenDay() { let q = make(.claude, claude: claudeSnapshot(five: 40, seven: 12)) XCTAssertEqual(q.short?.utilization, 40) diff --git a/notify.conf.example b/notify.conf.example index d5d00f8..897c1b4 100644 --- a/notify.conf.example +++ b/notify.conf.example @@ -74,6 +74,16 @@ # Default: tag #STACKNUDGE_SUBAGENT_NUDGES=off +# Fair-usage budget for pi, in tokens per day, per lane. pi enforces no quota +# of its own — API models bill per token against your keys, local models are +# free — so these are your own allowances and nothing cuts you off at them. The +# Usage tab bars today and this week (seven days' worth) against them and names +# the number on every row. 0 drops that lane's rows. Also settable in +# Settings → Usage, which writes back here. +# Defaults: 1500000 API, 500000 local +#STACKNUDGE_PI_API_BUDGET=1500000 +#STACKNUDGE_PI_LOCAL_BUDGET=500000 + # Log a debug line (to the launchd log) explaining voice decisions — useful # when a notification played silently and you want to know why. # Default: false diff --git a/panel/Panel.swift b/panel/Panel.swift index 127a479..3830553 100644 --- a/panel/Panel.swift +++ b/panel/Panel.swift @@ -1111,6 +1111,9 @@ final class PanelController: NSObject, NSApplicationDelegate, PanelKeyDelegate, private let claudeCliQuotaProbe = ClaudeCliQuotaProbe() private let codexQuotaProbe = CodexQuotaProbe() private let antigravityUsageProbe = AntigravityUsageProbe() + // Lazy because it shares the nav's history store rather than opening its + // own, so pi's transcripts are parsed once for the graph and the budget. + private lazy var piUsageProbe = PiUsageProbe(store: nav.usageStore) private var quotaTimer: Timer? // Last outcome derived per repo+branch, alongside the git values it was // derived from, so refreshOutcomes can skip re-deriving what hasn't moved. @@ -2100,6 +2103,16 @@ final class PanelController: NSObject, NSApplicationDelegate, PanelKeyDelegate, self.nav.quotaErrors[.antigravity] = nil } } + // Pi budget — read from pi's own transcripts, no network and no CLI. + // Unlike the three above this is not a provider quota: the denominators + // are the user's own (see PiBudget), so there is no failure to surface + // either. nil means no pi usage in either window. + piUsageProbe.fetch(budget: nav.piBudget) { [weak self] snapshot in + guard let self, let snapshot else { return } + self.nav.piQuota = snapshot + self.nav.quotaLastUpdated = Date() + self.nav.quotaUpdatedAt[.pi] = Date() + } } // Public hook for the Usage tab's "Sync now" keystroke. diff --git a/panel/PanelNav.swift b/panel/PanelNav.swift index c1dc3b7..d250f0e 100644 --- a/panel/PanelNav.swift +++ b/panel/PanelNav.swift @@ -172,6 +172,7 @@ enum SettingsRow: Hashable { case soundEnabled, agentDoneSound, permissionSound case voiceEnabled, voice, voiceSpeed, speakHotkey, downloadVoiceModel case quotaTracking, quotaAlerts, alertThreshold, pollFrequency, contextAlert, showRemaining + case piApiBudget, piLocalBudget case githubLinks, hideShipped, disconnectGithub case historyPerSession case editPhrases, checkPermissions, openConfig, releaseNotes, checkUpdates, uninstall, quit @@ -211,7 +212,7 @@ extension SettingsRow: CaseIterable { .soundEnabled, .agentDoneSound, .permissionSound, .voiceEnabled, .voice, .voiceSpeed, .speakHotkey, .downloadVoiceModel, .quotaTracking, .quotaAlerts, .alertThreshold, .pollFrequency, - .contextAlert, .showRemaining, + .contextAlert, .showRemaining, .piApiBudget, .piLocalBudget, .githubLinks, .hideShipped, .disconnectGithub, .historyPerSession, .editPhrases, .checkPermissions, .openConfig, .releaseNotes, @@ -424,6 +425,21 @@ final class PanelNav: ObservableObject { // Antigravity (agy) usage from the running CLI's loopback RPC, populated by // AntigravityUsageProbe — the agy analogue of `quota`/`codexQuota`. @Published var antigravityQuota: AntigravityQuotaSnapshot? { didSet { widgetQuotaCache = nil } } + // Pi's self-imposed budget, populated by PiUsageProbe from pi's own + // transcripts. Not a provider quota — see PiBudget. + @Published var piQuota: PiQuotaSnapshot? { didSet { widgetQuotaCache = nil } } + @Published var piApiBudgetDaily: Int = PiBudget.apiDailyDefault + @Published var piLocalBudgetDaily: Int = PiBudget.localDailyDefault + + var piBudget: PiBudget { + PiBudget(apiDaily: piApiBudgetDaily, localDaily: piLocalBudgetDaily) + } + + static func stepBudget(_ current: Int, forward: Bool) -> Int { + let list = PiBudget.dailyOptions + let index = list.firstIndex(of: current) ?? 0 + return list[forward ? (index + 1) % list.count : (index - 1 + list.count) % list.count] + } // Bumped by PanelController after a handoff is upserted into the ledger so // the Tickets tab (OutcomesView) and its tab-strip count re-read the // in-memory HandoffLedger and reflect the new session live. The ledger @@ -673,7 +689,10 @@ final class PanelNav: ObservableObject { @Published var usageWindow: UsageWindow = .widest // Long-lived parse cache. Owned per-nav rather than global so nothing leaks // between instances, and so tests can drive a clean one. - private let usageStore = UsageHistoryStore() + // Shared with PiUsageProbe so pi's transcripts are parsed once for both the + // graph and the budget. The store retains the widest span any caller asks + // for, so the graph's 24h refresh can't evict the budget's week. + let usageStore = UsageHistoryStore() // Replayed transcript history for the selected client, keyed so switching // client doesn't show another client's numbers while a scan is in flight. @Published var usageSeries: UsageSeries? @@ -699,6 +718,7 @@ final class PanelNav: ObservableObject { case .claude: return quota?.hasTier == true case .codex: return codexQuota?.hasTier == true case .antigravity: return antigravityQuota?.hasTier == true + case .pi: return piQuota?.hasTier == true } } @@ -747,7 +767,8 @@ final class PanelNav: ObservableObject { let value = WidgetQuota.make(client: selectedUsageClient, claude: quota, codex: codexQuota, - antigravity: antigravityQuota) + antigravity: antigravityQuota, + pi: piQuota) widgetQuotaCache = value return value } @@ -834,7 +855,7 @@ final class PanelNav: ObservableObject { // pressing W during an in-flight scan would otherwise have its // re-bucket silently reverted by this completion, leaving the header // claiming one window while the chart showed another. - store.refresh(source: source, retaining: .widest) + store.refresh(source: source, retaining: UsageWindow.widest.seconds) DispatchQueue.main.async { [weak self] in guard let self else { return } // Pure computation over entries already in memory (~1 ms), so @@ -1221,7 +1242,8 @@ final class PanelNav: ObservableObject { .widgetContent, .mascot, .theme] case .usage: return [.quotaTracking, .quotaAlerts, .alertThreshold, - .pollFrequency, .contextAlert, .showRemaining] + .pollFrequency, .contextAlert, .showRemaining, + .piApiBudget, .piLocalBudget] case .integrations: return [.slackPaste, .slackIdentity, .slackTest, .slackEnabled, .slackIdle, .slackDetail, .slackStop, @@ -1385,6 +1407,12 @@ final class PanelNav: ObservableObject { // Same coercion for poll interval — snap to nearest valid option. let rawPoll = Int(config["STACKNUDGE_USAGE_POLL_MIN"] ?? "") ?? 5 quotaPollMinutes = Self.quotaPollMinuteOptions.min(by: { abs($0 - rawPoll) < abs($1 - rawPoll) }) ?? 5 + let rawPiApi = Int(config["STACKNUDGE_PI_API_BUDGET"] ?? "") ?? PiBudget.apiDailyDefault + piApiBudgetDaily = PiBudget.dailyOptions.min(by: { abs($0 - rawPiApi) < abs($1 - rawPiApi) }) + ?? PiBudget.apiDailyDefault + let rawPiLocal = Int(config["STACKNUDGE_PI_LOCAL_BUDGET"] ?? "") ?? PiBudget.localDailyDefault + piLocalBudgetDaily = PiBudget.dailyOptions.min(by: { abs($0 - rawPiLocal) < abs($1 - rawPiLocal) }) + ?? PiBudget.localDailyDefault let rawCtx = Int(config["STACKNUDGE_CONTEXT_ALERT_THRESHOLD"] ?? "") ?? 0 contextAlertThresholdK = Self.contextAlertThresholdOptions.min(by: { abs($0 - rawCtx) < abs($1 - rawCtx) }) ?? 0 eventHistoryEnabled = ConfigFile.bool(config, "STACKNUDGE_EVENT_HISTORY", default: true) @@ -1699,7 +1727,7 @@ final class PanelNav: ObservableObject { .soundEnabled, .agentDoneSound, .permissionSound, .voiceEnabled, .voice, .voiceSpeed, .downloadVoiceModel, .quotaTracking, .quotaAlerts, .alertThreshold, .pollFrequency, - .contextAlert, .showRemaining, + .contextAlert, .showRemaining, .piApiBudget, .piLocalBudget, .githubLinks, .hideShipped, .historyPerSession, .eventHistory, .slackEnabled, .slackIdle, .slackDetail, .slackStop: @@ -1904,6 +1932,12 @@ final class PanelNav: ObservableObject { let next = forward ? (idx + 1) % list.count : (idx - 1 + list.count) % list.count quotaAlertThreshold = list[next] ConfigFile.write(key: "STACKNUDGE_QUOTA_THRESHOLD", value: String(quotaAlertThreshold)) + case .piApiBudget: + piApiBudgetDaily = Self.stepBudget(piApiBudgetDaily, forward: forward) + ConfigFile.write(key: "STACKNUDGE_PI_API_BUDGET", value: String(piApiBudgetDaily)) + case .piLocalBudget: + piLocalBudgetDaily = Self.stepBudget(piLocalBudgetDaily, forward: forward) + ConfigFile.write(key: "STACKNUDGE_PI_LOCAL_BUDGET", value: String(piLocalBudgetDaily)) case .pollFrequency: let list = Self.quotaPollMinuteOptions let idx = list.firstIndex(of: quotaPollMinutes) ?? 2 diff --git a/panel/PiUsage.swift b/panel/PiUsage.swift new file mode 100644 index 0000000..09d0825 --- /dev/null +++ b/panel/PiUsage.swift @@ -0,0 +1,133 @@ +import Foundation + +// Pi enforces no quota: its API models bill per token against the user's own +// keys, and its local models cost nothing. These denominators are therefore the +// user's own (STACKNUDGE_PI_API_BUDGET / STACKNUDGE_PI_LOCAL_BUDGET), not +// anything that will cut them off, which is why the pane says "budget" +// throughout — a bar that looks like Claude's would be read as a limit. +struct PiBudget: Equatable { + // Tokens per day, excluding cache reads. + var apiDaily: Int + var localDaily: Int + + // A week is seven days' worth. One knob per lane keeps the pair consistent: + // a separately-set weekly cap can sit below the daily one. + var apiWeekly: Int { apiDaily * 7 } + var localWeekly: Int { localDaily * 7 } + + static let apiDailyDefault = 1_500_000 + static let localDailyDefault = 500_000 + + static let fallback = PiBudget(apiDaily: apiDailyDefault, localDaily: localDailyDefault) + + // Ladder for the Settings cycle rows. 0 turns a lane off: someone who only + // runs local models has no use for an API budget, and a row pinned at 0% + // forever is worse than no row. + static let dailyOptions: [Int] = [0, 250_000, 500_000, 1_000_000, 1_500_000, + 2_000_000, 3_000_000, 5_000_000, 10_000_000] + + static func label(_ tokens: Int) -> String { + tokens == 0 ? "Off" : TokenFormat.short(tokens) + } +} + +// Pi's budget as the Usage tab's shared tier shape. Tiers are nil where the lane +// has no usage in the window, so a local-only user never sees an empty API row. +struct PiQuotaSnapshot: Equatable { + let apiToday: QuotaTier? + let apiThisWeek: QuotaTier? + let localToday: QuotaTier? + let localThisWeek: QuotaTier? + let budget: PiBudget + + var hasTier: Bool { + apiToday != nil || apiThisWeek != nil || localToday != nil || localThisWeek != nil + } + + // Sits where the other clients show their subscription tier. Naming a plan + // Pi doesn't have would be the one claim this pane must not make. + var planType: String? { "budget" } +} + +enum PiUsageBudget { + + // Calendar windows, not trailing ones: a real boundary is what gives the + // pane's "Resets" line and its ahead-of-pace warning something to measure + // against. dateInterval also reports the period's true length, so a 23-hour + // DST day paces correctly, and honours the locale's first weekday. + static func windows(now: Date, calendar: Calendar = .current) -> (day: DateInterval, week: DateInterval)? { + guard let day = calendar.dateInterval(of: .day, for: now), + let week = calendar.dateInterval(of: .weekOfYear, for: now) + else { return nil } + return (day, week) + } + + // Far enough back for the weekly window to be complete from the first poll. + static let retention: TimeInterval = 8 * 86400 + + static func snapshot(day: DateInterval, + dayTotals: UsageTotals, + week: DateInterval, + weekTotals: UsageTotals, + budget: PiBudget) -> PiQuotaSnapshot? { + let snapshot = PiQuotaSnapshot( + apiToday: tier(tokens: dayTotals.apiTokens, budget: budget.apiDaily, window: day), + apiThisWeek: tier(tokens: weekTotals.apiTokens, budget: budget.apiWeekly, window: week), + localToday: tier(tokens: dayTotals.localTokens, budget: budget.localDaily, window: day), + localThisWeek: tier(tokens: weekTotals.localTokens, budget: budget.localWeekly, window: week), + budget: budget) + return snapshot.hasTier ? snapshot : nil + } + + // Left unclamped deliberately: going over a self-imposed budget is the one + // thing it exists to tell you, and the bar clamps its own width anyway. + private static func tier(tokens: Int, budget: Int, window: DateInterval) -> QuotaTier? { + guard tokens > 0, budget > 0 else { return nil } + return QuotaTier(utilization: Double(tokens) / Double(budget) * 100, + resetsAt: window.end, + windowLength: window.duration) + } +} + +// Reads pi's transcripts through the same store the Usage graph uses, so a poll +// costs the directory walk rather than a re-parse. +final class PiUsageProbe { + + private let store: UsageHistoryStore + private let root: String? + // Serialises refreshes so an overlapping poll can't re-enter the store's + // scan, matching CodexQuotaProbe's probeQueue. + private let probeQueue = DispatchQueue(label: "stack-nudge.pi-usage") + + init(store: UsageHistoryStore, root: String? = nil) { + self.store = store + self.root = root + } + + // Calls completion on the main queue. nil means nothing to show — no pi + // usage in either window, or both lanes budgeted off. + func fetch(budget: PiBudget, + now: Date = Date(), + calendar: Calendar = .current, + completion: @escaping (PiQuotaSnapshot?) -> Void) { + probeQueue.async { [store, root] in + let result = Self.read(store: store, root: root, budget: budget, now: now, calendar: calendar) + DispatchQueue.main.async { completion(result) } + } + } + + static func read(store: UsageHistoryStore, + root: String?, + budget: PiBudget, + now: Date, + calendar: Calendar) -> PiQuotaSnapshot? { + guard let windows = PiUsageBudget.windows(now: now, calendar: calendar) else { return nil } + store.refresh(source: .pi, root: root, now: now, retaining: PiUsageBudget.retention) + return PiUsageBudget.snapshot( + day: windows.day, + dayTotals: store.totals(source: .pi, from: windows.day.start, to: windows.day.end), + week: windows.week, + weekTotals: store.totals(source: .pi, from: windows.week.start, to: windows.week.end), + budget: budget) + } +} diff --git a/panel/SessionUsage.swift b/panel/SessionUsage.swift index c5033cb..8cc2184 100644 --- a/panel/SessionUsage.swift +++ b/panel/SessionUsage.swift @@ -51,17 +51,19 @@ enum UsageClient: String, CaseIterable, Hashable { case claude case codex case antigravity + case pi var displayName: String { switch self { case .claude: return "Claude" case .codex: return "Codex" case .antigravity: return "Antigravity" + case .pi: return "Pi" } } // Where this client's replayable token history lives, or nil when it has - // none. Claude and Codex both write per-turn token usage with a timestamp + // none. Claude, Codex and pi all write per-turn token usage with a timestamp // into their transcripts. Antigravity's history.jsonl carries only prompt // text, timestamp and workspace — no token data at all — and its quota comes // from a live API call, so there's nothing to plot retrospectively. @@ -73,6 +75,7 @@ enum UsageClient: String, CaseIterable, Hashable { case .claude: return .claude case .codex: return .codex case .antigravity: return nil + case .pi: return .pi } } @@ -272,6 +275,7 @@ struct UsageView: View { case .claude: return nav.quota?.hasTier ?? false case .codex: return nav.codexQuota?.hasTier ?? false case .antigravity: return nav.antigravityQuota?.hasTier ?? false + case .pi: return nav.piQuota?.hasTier ?? false } } @@ -388,6 +392,24 @@ struct UsageView: View { section("Credits") { creditsRow(agy) } } } + case .pi: + if let pi = nav.piQuota { + // Every row names its denominator: these are the user's own + // budgets, not a limit pi will enforce, and a bare percentage + // here would read like Claude's above it. + if let tier = pi.apiToday { + section("API models today") { tierRow(tier, budget: pi.budget.apiDaily) } + } + if let tier = pi.apiThisWeek { + section("API models this week") { tierRow(tier, budget: pi.budget.apiWeekly) } + } + if let tier = pi.localToday { + section("Local models today") { tierRow(tier, budget: pi.budget.localDaily) } + } + if let tier = pi.localThisWeek { + section("Local models this week") { tierRow(tier, budget: pi.budget.localWeekly) } + } + } } } @@ -534,6 +556,7 @@ struct UsageView: View { case .claude: return nav.quota?.planType?.capitalized case .codex: return nav.codexQuota?.planType?.capitalized case .antigravity: return nav.antigravityQuota?.planType?.capitalized + case .pi: return nav.piQuota?.planType?.capitalized } } @@ -548,7 +571,7 @@ struct UsageView: View { } } - private func tierRow(_ tier: QuotaTier) -> some View { + private func tierRow(_ tier: QuotaTier, budget: Int? = nil) -> some View { // Show "30% used" or "70% remaining" depending on the toggle. Bar // still represents utilization so the color ramp keeps its meaning. let display = nav.quotaShowRemaining @@ -558,6 +581,11 @@ struct UsageView: View { return VStack(alignment: .leading, spacing: 4) { HStack(alignment: .firstTextBaseline, spacing: 8) { Spacer() + if let budget { + Text("of \(TokenFormat.short(budget)) budget") + .font(.caption2) + .foregroundStyle(.tertiary) + } Text("\(Int(display.rounded()))\(suffix)") .font(.caption.monospacedDigit().weight(.semibold)) .foregroundStyle(barColor(tier.utilization)) diff --git a/panel/Settings.swift b/panel/Settings.swift index 582fa86..b6449f5 100644 --- a/panel/Settings.swift +++ b/panel/Settings.swift @@ -627,6 +627,10 @@ struct SettingsView: View { case .pollFrequency: row(.pollFrequency, label: "Poll frequency", kind: .cycle, value: "\(nav.quotaPollMinutes) min", enabled: nav.quotaTrackingEnabled) case .contextAlert: row(.contextAlert, label: "Context alert at", kind: .cycle, value: contextAlertLabel) case .showRemaining: row(.showRemaining, label: "Show remaining", kind: .toggle, value: nav.quotaShowRemaining ? "On" : "Off", enabled: nav.quotaTrackingEnabled) + // Pi enforces no quota, so these are the user's own daily allowances and + // the Usage tab bars them against nothing else. The week is seven of them. + case .piApiBudget: row(.piApiBudget, label: "Pi API budget / day", kind: .cycle, value: PiBudget.label(nav.piApiBudgetDaily), enabled: nav.quotaTrackingEnabled) + case .piLocalBudget: row(.piLocalBudget, label: "Pi local budget / day", kind: .cycle, value: PiBudget.label(nav.piLocalBudgetDaily), enabled: nav.quotaTrackingEnabled) // Integrations case .slackPaste: row(.slackPaste, label: "Paste Slack setup", kind: .action, value: slackPasteValue) diff --git a/panel/UsageHistory.swift b/panel/UsageHistory.swift index 4de1921..940251c 100644 --- a/panel/UsageHistory.swift +++ b/panel/UsageHistory.swift @@ -78,9 +78,12 @@ enum UsageWindow: CaseIterable { var label: String { "\(hours)h" } - // The widest window, and therefore how much history the store retains. Every - // shorter window is bucketed from the same cached entries, so switching - // scope costs no I/O at all. + var seconds: TimeInterval { TimeInterval(hours * 3600) } + + // The widest window the graph plots. Every shorter one is bucketed from the + // same cached entries, so switching scope costs no I/O at all. It is not a + // ceiling on what the store keeps: retention is per source and per caller, + // and Pi's budget asks for a calendar week. static var widest: UsageWindow { .twentyFourHours } } @@ -114,6 +117,15 @@ struct UsageSeries: Equatable { } } +// Local/API token split over an arbitrary span. The budget's windows are +// calendar days and weeks, which no bucket grid lines up with and which run +// wider than any UsageWindow, so it sums entries rather than buckets. +struct UsageTotals: Equatable { + var localTokens = 0 + var apiTokens = 0 + var turns = 0 +} + // MARK: - Parsed entry // One agent turn, kept in the store so re-bucketing for a different window — or @@ -128,6 +140,16 @@ struct UsageEntry: Equatable { let outputTokens: Int let cacheReadTokens: Int let cacheCreateTokens: Int + // Whether the turn ran on a model that costs nothing. Only pi has one, so + // Claude and Codex entries are always billed. + let isLocal: Bool +} + +extension UsageEntry { + // The measure the graph's `excludingCacheReads` metric plots. Cache reads + // are replays of content already paid for, so counting them would let a + // resumed session burn a budget without doing any new work. + var tokensExcludingCacheReads: Int { inputTokens + outputTokens + cacheCreateTokens } } // MARK: - Store @@ -148,11 +170,13 @@ final class UsageHistoryStore { enum Source { case claude case codex + case pi var defaultRoot: String { switch self { case .claude: return "\(NSHomeDirectory())/.claude/projects" case .codex: return "\(NSHomeDirectory())/.codex/sessions" + case .pi: return "\(NSHomeDirectory())/.pi/agent/sessions" } } @@ -162,6 +186,7 @@ final class UsageHistoryStore { switch self { case .claude: return "\"assistant\"" case .codex: return "token_count" + case .pi: return "\"assistant\"" } } } @@ -189,6 +214,11 @@ final class UsageHistoryStore { // client, and a series must never merge another agent's entries. Flattening // one shared map here silently showed Claude's history under Codex. private var cache: [Source: [String: CachedFile]] = [:] + // Widest retention any caller has asked for, per source. Two callers read + // the same source over different spans — the graph plots 24h, the Pi budget + // sums a calendar week — and a narrow refresh must not evict the wide + // caller's entries out from under it, so the cutoff is the widest of them. + private var retention: [Source: TimeInterval] = [:] private var stats = RefreshStats() // Guards `cache` and `stats` only. File I/O and parsing happen outside it, so // a background refresh never blocks a main-thread re-bucket for a scope @@ -202,15 +232,19 @@ final class UsageHistoryStore { // MARK: Refresh - // Bring the cache up to date for `source`, retaining `retaining` hours of - // history. Safe to call repeatedly; the cost of a call that finds nothing - // changed is the directory walk plus a stat per file. + // Bring the cache up to date for `source`, retaining at least `retaining` + // seconds of history. Safe to call repeatedly; the cost of a call that finds + // nothing changed is the directory walk plus a stat per file. func refresh(source: Source, root: String? = nil, now: Date = Date(), - retaining: UsageWindow = .widest) { + retaining: TimeInterval = UsageWindow.widest.seconds) { let directory = root ?? source.defaultRoot - let cutoff = now.addingTimeInterval(-Double(retaining.hours * 3600)) + lock.lock() + let retained = max(retention[source] ?? 0, retaining) + retention[source] = retained + lock.unlock() + let cutoff = now.addingTimeInterval(-retained) var fresh = RefreshStats() for candidate in Self.transcriptFiles(under: directory, modifiedAfter: cutoff) { @@ -286,6 +320,30 @@ final class UsageHistoryStore { return accumulator.series() } + // Sum one source's cached entries over [from, to), split by local and API. + // Deduplicated by identity the same way bucketing is, and pure computation + // over what refresh already cached — the caller is responsible for having + // retained a span this wide. + func totals(source: Source, from: Date, to: Date) -> UsageTotals { + lock.lock() + let all = (cache[source] ?? [:]).values.flatMap { $0.entries } + lock.unlock() + + var seen: Set = [] + var result = UsageTotals() + for entry in all { + if let identity = entry.identity, !seen.insert(identity).inserted { continue } + guard entry.when >= from, entry.when < to else { continue } + result.turns += 1 + if entry.isLocal { + result.localTokens += entry.tokensExcludingCacheReads + } else { + result.apiTokens += entry.tokensExcludingCacheReads + } + } + return result + } + // Refresh then bucket, for callers that want both in one step. func refreshedSeries(source: Source, root: String? = nil, @@ -294,7 +352,7 @@ final class UsageHistoryStore { bucketSeconds: Int = UsageHistoryStore.defaultBucketSeconds) -> UsageSeries { // Always retain the widest window so narrower ones are a re-bucket of the // same cached entries rather than another scan. - refresh(source: source, root: root, now: now, retaining: .widest) + refresh(source: source, root: root, now: now, retaining: UsageWindow.widest.seconds) return series(source: source, now: now, window: window, bucketSeconds: bucketSeconds) } @@ -362,9 +420,43 @@ final class UsageHistoryStore { switch source { case .claude: return parseClaude(line: line) case .codex: return parseCodex(line: line) + case .pi: return parsePi(line: line) } } + private static func parsePi(line: Data) -> UsageEntry? { + guard let object = try? JSONSerialization.jsonObject(with: line) as? [String: Any], + (object["type"] as? String) == "message", + let rawTimestamp = object["timestamp"] as? String, + let when = parseTimestamp(rawTimestamp), + let message = object["message"] as? [String: Any], + (message["role"] as? String) == "assistant", + let usage = message["usage"] as? [String: Any] + else { return nil } + + let input = usage["input"] as? Int ?? 0 + let output = usage["output"] as? Int ?? 0 + let cacheRead = usage["cacheRead"] as? Int ?? 0 + let cacheWrite = usage["cacheWrite"] as? Int ?? 0 + let reasoning = usage["reasoning"] as? Int ?? 0 + + // pi writes the same usage block for a local model as for a billed one, + // so price is the only thing separating them. Read as NSNumber: a local + // turn's cost is the integer 0 and a billed one a fraction. + let cost = (usage["cost"] as? [String: Any]) + .flatMap { ($0["total"] as? NSNumber)?.doubleValue } ?? 0 + + return UsageEntry( + when: when, + identity: object["id"] as? String, + inputTokens: input, + outputTokens: output + reasoning, + cacheReadTokens: cacheRead, + cacheCreateTokens: cacheWrite, + isLocal: cost <= 0 + ) + } + private static func parseClaude(line: Data) -> UsageEntry? { guard let object = try? JSONSerialization.jsonObject(with: line) as? [String: Any], (object["type"] as? String) == "assistant", @@ -379,7 +471,8 @@ final class UsageHistoryStore { inputTokens: usage["input_tokens"] as? Int ?? 0, outputTokens: usage["output_tokens"] as? Int ?? 0, cacheReadTokens: usage["cache_read_input_tokens"] as? Int ?? 0, - cacheCreateTokens: usage["cache_creation_input_tokens"] as? Int ?? 0 + cacheCreateTokens: usage["cache_creation_input_tokens"] as? Int ?? 0, + isLocal: false ) } @@ -408,7 +501,8 @@ final class UsageHistoryStore { inputTokens: max(0, rawInput - cached), outputTokens: output + reasoning, cacheReadTokens: cached, - cacheCreateTokens: 0 + cacheCreateTokens: 0, + isLocal: false ) } diff --git a/panel/WidgetQuota.swift b/panel/WidgetQuota.swift index 1b5188d..73ba037 100644 --- a/panel/WidgetQuota.swift +++ b/panel/WidgetQuota.swift @@ -45,6 +45,8 @@ struct WidgetQuota: Equatable { switch client { case .antigravity: return "Inner ring: model closest to its limit · Outer ring: monthly prompt credits" + case .pi: + return "Inner ring: today's budget · Outer ring: this week's budget" default: let inner = "Inner ring: \(Self.ringPhrase(shortLabel, short))" guard long != nil else { return inner } @@ -74,7 +76,8 @@ struct WidgetQuota: Equatable { static func make(client: UsageClient?, claude: QuotaSnapshot?, codex: CodexQuotaSnapshot?, - antigravity: AntigravityQuotaSnapshot?) -> WidgetQuota { + antigravity: AntigravityQuotaSnapshot?, + pi: PiQuotaSnapshot?) -> WidgetQuota { switch client { case .claude: return WidgetQuota(client: .claude, @@ -102,6 +105,16 @@ struct WidgetQuota: Equatable { short: worst?.tier, long: creditsTier(antigravity?.promptCredits), shortLabel: "now", longLabel: "mo") + case .pi: + // Both rings read against the user's own budget. API takes the ring + // whenever it has usage: local tokens are free, so an API overrun is + // the one worth a glance. Labels are named rather than measured from + // the window — a DST day is 23 hours long and "23h" in a 35pt legend + // slot would be noise. + return WidgetQuota(client: .pi, + short: pi?.apiToday ?? pi?.localToday, + long: pi?.apiThisWeek ?? pi?.localThisWeek, + shortLabel: "1d", longLabel: "7d") case nil: return .empty } @@ -126,6 +139,7 @@ extension UsageClient { case .claude: return "Claude" case .codex: return "Codex" case .antigravity: return "Agy" + case .pi: return "Pi" } } } From faac591aed49a37fb403eab63432065d8257687f Mon Sep 17 00:00:00 2001 From: StuBehan Date: Wed, 23 Sep 2026 13:57:17 +0100 Subject: [PATCH 2/3] fix(usage): re-read pi usage when a budget changes --- .../PiUsageTests.swift | 43 +++++++++++++++++++ panel/Panel.swift | 19 ++++---- panel/PanelNav.swift | 25 ++++++++++- 3 files changed, 77 insertions(+), 10 deletions(-) diff --git a/Tests/StackNudgePanelCoreTests/PiUsageTests.swift b/Tests/StackNudgePanelCoreTests/PiUsageTests.swift index ff734d2..5f61823 100644 --- a/Tests/StackNudgePanelCoreTests/PiUsageTests.swift +++ b/Tests/StackNudgePanelCoreTests/PiUsageTests.swift @@ -285,3 +285,46 @@ final class PiUsageTests: XCTestCase { XCTAssertEqual(windows?.week.duration, 7 * 86400) } } + +// A budget change in Settings has to repaint the Usage tab straight away, not +// on the next poll, and a lane switched off has to take its row with it. +@MainActor +final class PiBudgetSettingsTests: XCTestCase { + + private func snapshot() -> PiQuotaSnapshot { + let day = DateInterval(start: Date(), duration: 86400) + return PiQuotaSnapshot(apiToday: QuotaTier(utilization: 40, resetsAt: day.end, windowLength: day.duration), + apiThisWeek: nil, localToday: nil, localThisWeek: nil, + budget: .fallback) + } + + func test_budgetChangeRefreshesStraightAway() { + let nav = PanelNav() + var refreshes = 0 + nav.refreshPiBudget = { refreshes += 1 } + + nav.setPiBudget(apiDaily: 2_000_000, localDaily: 0) + + XCTAssertEqual(refreshes, 1) + XCTAssertEqual(nav.piBudget, PiBudget(apiDaily: 2_000_000, localDaily: 0)) + } + + func test_nilSnapshotClearsTheRow() { + let nav = PanelNav() + nav.applyPiSnapshot(snapshot()) + + nav.applyPiSnapshot(nil) + + XCTAssertNil(nav.piQuota) + XCTAssertFalse(nav.availableUsageClients.contains(.pi)) + } + + func test_snapshotStampsPiAsUpdated() { + let nav = PanelNav() + + nav.applyPiSnapshot(snapshot()) + + XCTAssertNotNil(nav.quotaUpdatedAt[.pi]) + XCTAssertTrue(nav.availableUsageClients.contains(.pi)) + } +} diff --git a/panel/Panel.swift b/panel/Panel.swift index 3830553..181c17b 100644 --- a/panel/Panel.swift +++ b/panel/Panel.swift @@ -1311,6 +1311,7 @@ final class PanelController: NSObject, NSApplicationDelegate, PanelKeyDelegate, nav.refreshOutcomes = { [weak self] in self?.refreshOutcomes() } nav.refreshPullRequests = { [weak self] in self?.refreshPullRequests() } nav.refreshPullRequestsNow = { [weak self] in self?.refreshPullRequestsNow() } + nav.refreshPiBudget = { [weak self] in self?.refreshPiUsage() } nav.startGithubSignIn = { [weak self] in self?.startGithubSignIn() } nav.cancelGithubSignIn = { [weak self] in self?.cancelGithubSignIn() } @@ -2103,15 +2104,17 @@ final class PanelController: NSObject, NSApplicationDelegate, PanelKeyDelegate, self.nav.quotaErrors[.antigravity] = nil } } - // Pi budget — read from pi's own transcripts, no network and no CLI. - // Unlike the three above this is not a provider quota: the denominators - // are the user's own (see PiBudget), so there is no failure to surface - // either. nil means no pi usage in either window. + refreshPiUsage() + } + + // Pi budget, read from pi's own transcripts with no network and no CLI. + // Unlike the probes above there's no failure to surface: the denominators + // are the user's own (see PiBudget). Also run on a budget change in + // Settings, which is why it sits outside runQuotaProbe. + private func refreshPiUsage() { + guard quotaTrackingEnabled else { return } piUsageProbe.fetch(budget: nav.piBudget) { [weak self] snapshot in - guard let self, let snapshot else { return } - self.nav.piQuota = snapshot - self.nav.quotaLastUpdated = Date() - self.nav.quotaUpdatedAt[.pi] = Date() + self?.nav.applyPiSnapshot(snapshot) } } diff --git a/panel/PanelNav.swift b/panel/PanelNav.swift index d250f0e..3cedc02 100644 --- a/panel/PanelNav.swift +++ b/panel/PanelNav.swift @@ -434,6 +434,25 @@ final class PanelNav: ObservableObject { var piBudget: PiBudget { PiBudget(apiDaily: piApiBudgetDaily, localDaily: piLocalBudgetDaily) } + // Wired by PanelController to re-read pi's usage. Fired on a budget change + // so the Usage tab isn't left on the old denominator until the next poll. + var refreshPiBudget: (() -> Void)? + + func setPiBudget(apiDaily: Int, localDaily: Int) { + piApiBudgetDaily = apiDaily + piLocalBudgetDaily = localDaily + refreshPiBudget?() + } + + // nil clears the row rather than holding the last snapshot. pi is read from + // local disk, so there's no dropped tick to ride out: nil means no usage in + // either window, or both lanes budgeted off. + func applyPiSnapshot(_ snapshot: PiQuotaSnapshot?) { + piQuota = snapshot + guard snapshot != nil else { return } + quotaLastUpdated = Date() + quotaUpdatedAt[.pi] = Date() + } static func stepBudget(_ current: Int, forward: Bool) -> Int { let list = PiBudget.dailyOptions @@ -1933,10 +1952,12 @@ final class PanelNav: ObservableObject { quotaAlertThreshold = list[next] ConfigFile.write(key: "STACKNUDGE_QUOTA_THRESHOLD", value: String(quotaAlertThreshold)) case .piApiBudget: - piApiBudgetDaily = Self.stepBudget(piApiBudgetDaily, forward: forward) + setPiBudget(apiDaily: Self.stepBudget(piApiBudgetDaily, forward: forward), + localDaily: piLocalBudgetDaily) ConfigFile.write(key: "STACKNUDGE_PI_API_BUDGET", value: String(piApiBudgetDaily)) case .piLocalBudget: - piLocalBudgetDaily = Self.stepBudget(piLocalBudgetDaily, forward: forward) + setPiBudget(apiDaily: piApiBudgetDaily, + localDaily: Self.stepBudget(piLocalBudgetDaily, forward: forward)) ConfigFile.write(key: "STACKNUDGE_PI_LOCAL_BUDGET", value: String(piLocalBudgetDaily)) case .pollFrequency: let list = Self.quotaPollMinuteOptions From 39d0d2fb87b4e3453a1669fc99d123ffcada5dde Mon Sep 17 00:00:00 2001 From: StuBehan Date: Wed, 23 Sep 2026 14:47:57 +0100 Subject: [PATCH 3/3] feat(usage): list pi usage per model with a today/week toggle --- README.md | 2 +- .../PiUsageTests.swift | 156 +++++++++++++----- .../WidgetQuotaTests.swift | 34 ++-- panel/Panel.swift | 3 + panel/PanelNav.swift | 9 + panel/PiUsage.swift | 80 +++++++-- panel/SessionUsage.swift | 78 +++++---- panel/UsageHistory.swift | 56 ++++--- panel/WidgetQuota.swift | 18 +- 9 files changed, 299 insertions(+), 137 deletions(-) diff --git a/README.md b/README.md index 7840b2f..994308d 100644 --- a/README.md +++ b/README.md @@ -261,7 +261,7 @@ Claude, Codex, Antigravity and pi each appear in the tab's client list when they For Claude and Codex the two rings are the 5-hour and weekly windows — though Codex reports the window length per limit rather than a fixed pair, and on some accounts publishes only a weekly one, so its Usage tab headings are named from what it actually reports. Antigravity reports neither — it publishes one window per model plus a monthly credit pool — so its inner ring shows whichever model is closest to its limit and the outer ring shows monthly prompt credits. Gemini CLI has no usage counter: unlike the others it writes no rate-limit data to disk and serves no local endpoint, so there's nothing to read. -**pi is the one client whose bars are not a quota.** It enforces none: its API models bill per token against your own keys, and its local models cost nothing. What its rows track is a fair-usage budget you set yourself — a daily token allowance per lane, with the week being seven of them — so every row names its denominator (*"of 1.5M budget"*) and the tab labels it `Budget` where the others show a subscription tier. The windows are calendar ones, today and this week, which is what makes the reset time and the ahead-of-pace warning mean anything. Tokens are counted the way the history graph's "no cache reads" metric counts them, so resuming a long session doesn't burn the budget by replaying itself, and a turn counts as local when pi priced it at zero. Set the two allowances in Settings → Usage (`STACKNUDGE_PI_API_BUDGET` / `STACKNUDGE_PI_LOCAL_BUDGET`, tokens per day, `Off` to drop a lane's rows). Going over is reported rather than clamped: that is the only thing a budget you set for yourself can usefully tell you. +**pi is the one client whose bars are not a quota.** It enforces none: its API models bill per token against your own keys, and its local models cost nothing. What its page tracks is a fair-usage budget you set yourself, one daily token allowance for API models and one for local models, with the week being seven days of it. The page lists one row per model, the same bar, reset line and pace warning as the other clients, each measured against the allowance for its kind and naming it (*"of 1.5M API budget"*), and the tab labels it `Budget` where the others show a subscription tier. `W` toggles the page between today and this week. The windows are calendar ones rather than rolling, which is what gives the reset time and the ahead-of-pace warning something to measure against, and the widget's rings follow whichever model is closest to its budget in each. Tokens are counted the way the history graph's "no cache reads" metric counts them, so resuming a long session doesn't burn the budget by replaying itself, and a turn counts as local when pi priced it at zero. Set the two allowances in Settings → Usage (`STACKNUDGE_PI_API_BUDGET` / `STACKNUDGE_PI_LOCAL_BUDGET`, tokens per day, `Off` to drop that kind's models from the page). Going over is reported rather than clamped: that is the only thing a budget you set for yourself can usefully tell you. #### Threshold-crossing notifications diff --git a/Tests/StackNudgePanelCoreTests/PiUsageTests.swift b/Tests/StackNudgePanelCoreTests/PiUsageTests.swift index 5f61823..2aeb08a 100644 --- a/Tests/StackNudgePanelCoreTests/PiUsageTests.swift +++ b/Tests/StackNudgePanelCoreTests/PiUsageTests.swift @@ -2,6 +2,12 @@ import XCTest @testable import StackNudgePanelCore +// Lane sums for asserting which kind a turn was filed under. +private extension Dictionary where Key == UsageModelKey, Value == Int { + var localTokens: Int { filter { $0.key.isLocal }.values.reduce(0, +) } + var apiTokens: Int { filter { !$0.key.isLocal }.values.reduce(0, +) } +} + // Pi's row on the Usage tab is a budget the user sets, not a quota a provider // enforces, so what these pin is that every number under it is derived rather // than assumed: which turns count as local, how wide the windows really are, @@ -40,16 +46,18 @@ final class PiUsageTests: XCTestCase { cacheWrite: Int = 0, reasoning: Int = 0, cost: Double = 0, + model: String = "m", + provider: String = "p", id: String = UUID().uuidString) -> String { """ {"type":"message","id":"\(id)","timestamp":"\(Self.stamp.string(from: when))",\ - "message":{"role":"assistant","model":"m","usage":{"input":\(input),\ + "message":{"role":"assistant","model":"\(model)","provider":"\(provider)","usage":{"input":\(input),\ "output":\(output),"cacheRead":\(cacheRead),"cacheWrite":\(cacheWrite),\ "reasoning":\(reasoning),"totalTokens":0,"cost":{"total":\(cost)}}}} """ } - private func totals(_ lines: [String], now: Date, span: TimeInterval = 8 * 86400) -> UsageTotals { + private func totals(_ lines: [String], now: Date, span: TimeInterval = 8 * 86400) -> [UsageModelKey: Int] { let dir = fixtureDirectory() write(lines, to: dir, name: "session.jsonl", modified: now) let store = UsageHistoryStore() @@ -111,7 +119,6 @@ final class PiUsageTests: XCTestCase { now: now) XCTAssertEqual(actual.apiTokens, 1_100) XCTAssertEqual(actual.localTokens, 1_100) - XCTAssertEqual(actual.turns, 2) } // A resumed session replays earlier turns into a second transcript. @@ -119,7 +126,6 @@ final class PiUsageTests: XCTestCase { let now = Date(timeIntervalSince1970: 1_785_150_000) let line = piLine(at: now.addingTimeInterval(-60), cost: 1, id: "repeated") let actual = totals([line, line], now: now) - XCTAssertEqual(actual.turns, 1) XCTAssertEqual(actual.apiTokens, 1_100) } @@ -131,7 +137,29 @@ final class PiUsageTests: XCTestCase { "message":{"role":"user","content":[{"type":"text","text":"ask the assistant"}]}} """ let actual = totals([session, user, piLine(at: now.addingTimeInterval(-60), cost: 1)], now: now) - XCTAssertEqual(actual.turns, 1) + XCTAssertEqual(actual.apiTokens, 1_100) + } + + // MARK: Per model + + func test_totalsSplitByModel() { + let now = Date(timeIntervalSince1970: 1_785_150_000) + let actual = totals([piLine(at: now.addingTimeInterval(-60), cost: 1, model: "gemini-pro", provider: "gemini"), + piLine(at: now.addingTimeInterval(-90), cost: 1, model: "gemini-pro", provider: "gemini"), + piLine(at: now.addingTimeInterval(-120), cost: 0, model: "qwen", provider: "ollama")], + now: now) + XCTAssertEqual(actual[UsageModelKey(provider: "gemini", model: "gemini-pro", isLocal: false)], 2_200) + XCTAssertEqual(actual[UsageModelKey(provider: "ollama", model: "qwen", isLocal: true)], 1_100) + } + + // A failed call is priced at zero and carries no tokens. It must not turn up + // as an empty local model row. + func test_zeroTokenTurnAddsNoModel() { + let now = Date(timeIntervalSince1970: 1_785_150_000) + let actual = totals([piLine(at: now.addingTimeInterval(-60), input: 0, output: 0, + model: "claude-3-5-haiku-latest", provider: "anthropic")], + now: now) + XCTAssertTrue(actual.isEmpty) } // MARK: Window arithmetic @@ -143,7 +171,6 @@ final class PiUsageTests: XCTestCase { let actual = totals([piLine(at: now.addingTimeInterval(-6 * 86400), cost: 1), piLine(at: now.addingTimeInterval(-60), cost: 1)], now: now) - XCTAssertEqual(actual.turns, 2) XCTAssertEqual(actual.apiTokens, 2_200) } @@ -157,7 +184,7 @@ final class PiUsageTests: XCTestCase { store.refresh(source: .pi, root: dir, now: now, retaining: 8 * 86400) let actual = store.totals(source: .pi, from: now.addingTimeInterval(-86400), to: now) - XCTAssertEqual(actual.turns, 1) + XCTAssertEqual(actual.apiTokens, 1_100) } // The graph refreshes the same source over 24h. Without a retention @@ -181,50 +208,81 @@ final class PiUsageTests: XCTestCase { DateInterval(start: Date(timeIntervalSince1970: start), duration: duration) } - func test_utilizationIsTokensOverBudget() { + private func modelTotals(_ rows: [(model: String, provider: String, isLocal: Bool, tokens: Int)]) -> [UsageModelKey: Int] { + var totals: [UsageModelKey: Int] = [:] + for row in rows { + totals[UsageModelKey(provider: row.provider, model: row.model, isLocal: row.isLocal)] = row.tokens + } + return totals + } + + private func today(_ totals: [UsageModelKey: Int], budget: PiBudget) -> [PiModelUsage] { let day = window(1_785_110_400, 86400) - let snapshot = PiUsageBudget.snapshot(day: day, - dayTotals: UsageTotals(localTokens: 0, apiTokens: 750_000, turns: 1), - week: day, - weekTotals: UsageTotals(localTokens: 0, apiTokens: 750_000, turns: 1), - budget: PiBudget(apiDaily: 1_500_000, localDaily: 500_000)) - XCTAssertEqual(snapshot?.apiToday?.utilization, 50) - XCTAssertEqual(snapshot?.apiToday?.resetsAt, day.end) - XCTAssertEqual(snapshot?.apiToday?.windowLength, 86400) - XCTAssertNil(snapshot?.localToday) + return PiUsageBudget.snapshot(day: day, dayTotals: totals, week: day, weekTotals: totals, + budget: budget)?.today ?? [] } - // Going over is the thing a self-imposed budget exists to report. - func test_overBudgetIsNotClamped() { + // The pane lists models by name only, but the kind still picks the budget: + // a local model against the local allowance, a priced one against the API's. + func test_eachModelMeasuresAgainstItsKindsBudget() { let day = window(1_785_110_400, 86400) - let snapshot = PiUsageBudget.snapshot(day: day, - dayTotals: UsageTotals(localTokens: 3_000_000, apiTokens: 0, turns: 1), - week: day, - weekTotals: UsageTotals(localTokens: 3_000_000, apiTokens: 0, turns: 1), - budget: PiBudget(apiDaily: 1_500_000, localDaily: 500_000)) - XCTAssertEqual(snapshot?.localToday?.utilization, 600) + let totals = modelTotals([("gemini-pro", "gemini", false, 750_000), + ("qwen3.8:27b", "ollama", true, 60_000)]) + let actual = today(totals, budget: PiBudget(apiDaily: 1_500_000, localDaily: 500_000)) + XCTAssertEqual(actual.map(\.name), ["gemini-pro", "qwen3.8:27b"]) + XCTAssertEqual(actual.map(\.tier.utilization), [50, 12]) + XCTAssertEqual(actual.first?.tier.resetsAt, day.end) + XCTAssertEqual(actual.first?.tier.windowLength, 86400) } - func test_weeklyBudgetIsSevenDays() { - XCTAssertEqual(PiBudget(apiDaily: 1_000_000, localDaily: 100_000).apiWeekly, 7_000_000) - XCTAssertEqual(PiBudget(apiDaily: 1_000_000, localDaily: 100_000).localWeekly, 700_000) + // Closest to its budget first, not most tokens: 400K local is nearer its + // limit than 750K of API. + func test_modelsSortClosestToBudgetFirst() { + let totals = modelTotals([("gemini-pro", "gemini", false, 750_000), + ("qwen3.8:27b", "ollama", true, 400_000)]) + let actual = today(totals, budget: PiBudget(apiDaily: 1_500_000, localDaily: 500_000)) + XCTAssertEqual(actual.map(\.name), ["qwen3.8:27b", "gemini-pro"]) } - // A lane switched off in Settings draws no row, however much it was used. - func test_zeroBudgetDrawsNoTier() { + func test_weekMeasuresAgainstSevenDaysOfAllowance() { let day = window(1_785_110_400, 86400) - let snapshot = PiUsageBudget.snapshot(day: day, - dayTotals: UsageTotals(localTokens: 400_000, apiTokens: 0, turns: 1), - week: day, - weekTotals: UsageTotals(localTokens: 400_000, apiTokens: 0, turns: 1), - budget: PiBudget(apiDaily: 1_500_000, localDaily: 0)) - XCTAssertNil(snapshot) + let week = window(1_785_110_400, 7 * 86400) + let totals = modelTotals([("gemini-pro", "gemini", false, 700_000)]) + let snapshot = PiUsageBudget.snapshot(day: day, dayTotals: [:], week: week, weekTotals: totals, + budget: PiBudget(apiDaily: 1_000_000, localDaily: 500_000)) + XCTAssertTrue(snapshot?.today.isEmpty ?? false) + XCTAssertEqual(snapshot?.thisWeek.first?.tier.utilization, 10) + XCTAssertEqual(snapshot?.thisWeek.first?.tier.resetsAt, week.end) + } + + func test_sameModelNameFromTwoProvidersNamesTheProvider() { + let totals = modelTotals([("qwen3", "qwen", false, 2_000), ("qwen3", "openrouter", false, 1_000)]) + XCTAssertEqual(today(totals, budget: .fallback).map(\.name), ["qwen3 (qwen)", "qwen3 (openrouter)"]) + } + + // A kind budgeted off in Settings takes its models off the page. + func test_kindBudgetedOffDropsItsModels() { + let totals = modelTotals([("gemini-pro", "gemini", false, 100_000), ("qwen", "ollama", true, 50_000)]) + XCTAssertEqual(today(totals, budget: PiBudget(apiDaily: 0, localDaily: 500_000)).map(\.name), ["qwen"]) + } + + // Going over is the thing a self-imposed budget exists to report. + func test_overBudgetIsNotClamped() { + let totals = modelTotals([("qwen", "ollama", true, 3_000_000)]) + XCTAssertEqual(today(totals, budget: PiBudget(apiDaily: 1_500_000, localDaily: 500_000)).first?.tier.utilization, 600) + } + + func test_weeklyBudgetIsSevenDays() { + let budget = PiBudget(apiDaily: 1_000_000, localDaily: 100_000) + XCTAssertEqual(budget.allowance(isLocal: false, in: .thisWeek), 7_000_000) + XCTAssertEqual(budget.allowance(isLocal: true, in: .thisWeek), 700_000) + XCTAssertEqual(budget.allowance(isLocal: true, in: .today), 100_000) } func test_noUsageProducesNoSnapshot() { let day = window(1_785_110_400, 86400) - XCTAssertNil(PiUsageBudget.snapshot(day: day, dayTotals: UsageTotals(), - week: day, weekTotals: UsageTotals(), + XCTAssertNil(PiUsageBudget.snapshot(day: day, dayTotals: [:], + week: day, weekTotals: [:], budget: .fallback)) } @@ -247,10 +305,10 @@ final class PiUsageTests: XCTestCase { now: now, calendar: calendar) - XCTAssertEqual(actual?.apiToday?.utilization, 0.1) - let week = try XCTUnwrap(actual?.apiThisWeek?.utilization) + XCTAssertEqual(actual?.today.map(\.name), ["m"]) + XCTAssertEqual(actual?.today.first?.tier.utilization, 0.1) + let week = try XCTUnwrap(actual?.thisWeek.first?.tier.utilization) XCTAssertEqual(week, 1_500.0 / 7_000_000 * 100, accuracy: 0.000_001) - XCTAssertNil(actual?.localToday) } // MARK: Calendar windows @@ -293,9 +351,10 @@ final class PiBudgetSettingsTests: XCTestCase { private func snapshot() -> PiQuotaSnapshot { let day = DateInterval(start: Date(), duration: 86400) - return PiQuotaSnapshot(apiToday: QuotaTier(utilization: 40, resetsAt: day.end, windowLength: day.duration), - apiThisWeek: nil, localToday: nil, localThisWeek: nil, - budget: .fallback) + let tier = QuotaTier(utilization: 40, resetsAt: day.end, windowLength: day.duration) + let model = PiModelUsage(key: UsageModelKey(provider: "gemini", model: "gemini-pro", isLocal: false), + name: "gemini-pro", tokens: 600_000, tier: tier) + return PiQuotaSnapshot(today: [model], thisWeek: [], budget: .fallback) } func test_budgetChangeRefreshesStraightAway() { @@ -309,6 +368,15 @@ final class PiBudgetSettingsTests: XCTestCase { XCTAssertEqual(nav.piBudget, PiBudget(apiDaily: 2_000_000, localDaily: 0)) } + func test_windowTogglesBetweenTodayAndThisWeek() { + let nav = PanelNav() + XCTAssertEqual(nav.piWindow, .today) + nav.cyclePiWindow() + XCTAssertEqual(nav.piWindow, .thisWeek) + nav.cyclePiWindow() + XCTAssertEqual(nav.piWindow, .today) + } + func test_nilSnapshotClearsTheRow() { let nav = PanelNav() nav.applyPiSnapshot(snapshot()) diff --git a/Tests/StackNudgePanelCoreTests/WidgetQuotaTests.swift b/Tests/StackNudgePanelCoreTests/WidgetQuotaTests.swift index 9ad4761..3194a7b 100644 --- a/Tests/StackNudgePanelCoreTests/WidgetQuotaTests.swift +++ b/Tests/StackNudgePanelCoreTests/WidgetQuotaTests.swift @@ -41,40 +41,38 @@ final class WidgetQuotaTests: XCTestCase { WidgetQuota.make(client: client, claude: claude, codex: codex, antigravity: agy, pi: pi) } - private func piSnapshot(apiToday: Double?, apiWeek: Double?, - localToday: Double?, localWeek: Double?) -> PiQuotaSnapshot { - let day = DateInterval(start: Date(), duration: 86400) - let week = DateInterval(start: Date(), duration: 7 * 86400) - func tier(_ used: Double?, _ window: DateInterval) -> QuotaTier? { - used.map { QuotaTier(utilization: $0, resetsAt: window.end, windowLength: window.duration) } + private func piSnapshot(today: [Double], week: [Double]) -> PiQuotaSnapshot { + func models(_ used: [Double], _ duration: TimeInterval) -> [PiModelUsage] { + used.enumerated().map { index, value in + PiModelUsage(key: UsageModelKey(provider: "p", model: "m\(index)", isLocal: false), + name: "m\(index)", tokens: 1, + tier: QuotaTier(utilization: value, resetsAt: Date().addingTimeInterval(duration), + windowLength: duration)) + } } - return PiQuotaSnapshot(apiToday: tier(apiToday, day), - apiThisWeek: tier(apiWeek, week), - localToday: tier(localToday, day), - localThisWeek: tier(localWeek, week), - budget: .fallback) + return PiQuotaSnapshot(today: models(today, 86400), thisWeek: models(week, 7 * 86400), budget: .fallback) } // MARK: - Per-client ring mapping - func test_pi_mapsTodayAndThisWeek() { - let q = make(.pi, pi: piSnapshot(apiToday: 62, apiWeek: 18, localToday: 4, localWeek: 2)) + // Like Antigravity: whichever model is closest to its budget takes the ring. + func test_pi_closestModelTakesEachRing() { + let q = make(.pi, pi: piSnapshot(today: [12, 62, 4], week: [18, 2])) XCTAssertEqual(q.short?.utilization, 62) XCTAssertEqual(q.long?.utilization, 18) XCTAssertEqual(q.shortLabel, "1d") XCTAssertEqual(q.longLabel, "7d") } - // Local-only usage still gets both rings rather than falling back to empty. - func test_pi_fallsBackToLocalWhenNoApiUsage() { - let q = make(.pi, pi: piSnapshot(apiToday: nil, apiWeek: nil, localToday: 30, localWeek: 9)) - XCTAssertEqual(q.short?.utilization, 30) + func test_pi_weekOnlyUsageLeavesTheInnerRingEmpty() { + let q = make(.pi, pi: piSnapshot(today: [], week: [9])) + XCTAssertNil(q.short) XCTAssertEqual(q.long?.utilization, 9) } // A budget is the user's own, so passing it is the point of the row. func test_pi_overBudgetIsNotClamped() { - let q = make(.pi, pi: piSnapshot(apiToday: 140, apiWeek: 30, localToday: nil, localWeek: nil)) + let q = make(.pi, pi: piSnapshot(today: [140], week: [30])) XCTAssertEqual(q.short?.utilization, 140) } diff --git a/panel/Panel.swift b/panel/Panel.swift index 181c17b..f2d739c 100644 --- a/panel/Panel.swift +++ b/panel/Panel.swift @@ -4227,6 +4227,9 @@ final class PanelController: NSObject, NSApplicationDelegate, PanelKeyDelegate, case KeyCode.wKey where nav.usagePane == .history: // Re-buckets cached entries; no rescan, so it lands instantly. nav.cycleUsageWindow() + case KeyCode.wKey where nav.selectedUsageClient == .pi: + // Both windows are already in the snapshot, so this is a repaint. + nav.cyclePiWindow() case KeyCode.rKey: syncQuotaNow() case KeyCode.pKey: diff --git a/panel/PanelNav.swift b/panel/PanelNav.swift index 3cedc02..96c74e6 100644 --- a/panel/PanelNav.swift +++ b/panel/PanelNav.swift @@ -434,6 +434,15 @@ final class PanelNav: ObservableObject { var piBudget: PiBudget { PiBudget(apiDaily: piApiBudgetDaily, localDaily: piLocalBudgetDaily) } + // Which window the pi page shows; W toggles it. In-memory, like usageWindow. + @Published var piWindow: PiWindow = .today + + func cyclePiWindow() { + let windows = PiWindow.allCases + let index = windows.firstIndex(of: piWindow) ?? 0 + piWindow = windows[(index + 1) % windows.count] + } + // Wired by PanelController to re-read pi's usage. Fired on a budget change // so the Usage tab isn't left on the old denominator until the next poll. var refreshPiBudget: (() -> Void)? diff --git a/panel/PiUsage.swift b/panel/PiUsage.swift index 09d0825..408b75f 100644 --- a/panel/PiUsage.swift +++ b/panel/PiUsage.swift @@ -29,21 +29,54 @@ struct PiBudget: Equatable { static func label(_ tokens: Int) -> String { tokens == 0 ? "Off" : TokenFormat.short(tokens) } + + func allowance(isLocal: Bool, in window: PiWindow) -> Int { + switch window { + case .today: return isLocal ? localDaily : apiDaily + case .thisWeek: return isLocal ? localWeekly : apiWeekly + } + } +} + +// One model over one window, measured against the budget for its kind: a local +// model against the local allowance, anything pi priced against the API one. +// The pane lists models by name only; the kind decides the denominator. +struct PiModelUsage: Equatable { + let key: UsageModelKey + let name: String + let tokens: Int + let tier: QuotaTier +} + +// Which window the pi page shows. W toggles it, like the History pane's window. +enum PiWindow: CaseIterable { + case today + case thisWeek + + var label: String { + switch self { + case .today: return "Today" + case .thisWeek: return "This week" + } + } } -// Pi's budget as the Usage tab's shared tier shape. Tiers are nil where the lane -// has no usage in the window, so a local-only user never sees an empty API row. +// Pi's budget as the Usage tab reads it. A model appears in a window only if it +// ran there and its kind has a budget, so a local-only user never sees an API row. struct PiQuotaSnapshot: Equatable { - let apiToday: QuotaTier? - let apiThisWeek: QuotaTier? - let localToday: QuotaTier? - let localThisWeek: QuotaTier? + let today: [PiModelUsage] + let thisWeek: [PiModelUsage] let budget: PiBudget - var hasTier: Bool { - apiToday != nil || apiThisWeek != nil || localToday != nil || localThisWeek != nil + func models(in window: PiWindow) -> [PiModelUsage] { + switch window { + case .today: return today + case .thisWeek: return thisWeek + } } + var hasTier: Bool { !today.isEmpty || !thisWeek.isEmpty } + // Sits where the other clients show their subscription tier. Naming a plan // Pi doesn't have would be the one claim this pane must not make. var planType: String? { "budget" } @@ -66,19 +99,38 @@ enum PiUsageBudget { static let retention: TimeInterval = 8 * 86400 static func snapshot(day: DateInterval, - dayTotals: UsageTotals, + dayTotals: [UsageModelKey: Int], week: DateInterval, - weekTotals: UsageTotals, + weekTotals: [UsageModelKey: Int], budget: PiBudget) -> PiQuotaSnapshot? { let snapshot = PiQuotaSnapshot( - apiToday: tier(tokens: dayTotals.apiTokens, budget: budget.apiDaily, window: day), - apiThisWeek: tier(tokens: weekTotals.apiTokens, budget: budget.apiWeekly, window: week), - localToday: tier(tokens: dayTotals.localTokens, budget: budget.localDaily, window: day), - localThisWeek: tier(tokens: weekTotals.localTokens, budget: budget.localWeekly, window: week), + today: models(dayTotals, interval: day, window: .today, budget: budget), + thisWeek: models(weekTotals, interval: week, window: .thisWeek, budget: budget), budget: budget) return snapshot.hasTier ? snapshot : nil } + // Closest to its budget first, which is the order the widget and a glance + // at the page both care about. + private static func models(_ totals: [UsageModelKey: Int], interval: DateInterval, + window: PiWindow, budget: PiBudget) -> [PiModelUsage] { + // The provider only earns a place in the name when two providers serve + // the same model name; otherwise it's noise in a narrow pane. + let shared = Set(Dictionary(grouping: totals.keys, by: \.model).filter { $0.value.count > 1 }.keys) + return totals.compactMap { key, used -> PiModelUsage? in + guard let tier = tier(tokens: used, budget: budget.allowance(isLocal: key.isLocal, in: window), + window: interval) + else { return nil } + let name = shared.contains(key.model) ? "\(key.model) (\(key.provider ?? "unknown"))" : key.model + return PiModelUsage(key: key, name: name, tokens: used, tier: tier) + } + .sorted { + $0.tier.utilization != $1.tier.utilization + ? $0.tier.utilization > $1.tier.utilization + : $0.name < $1.name + } + } + // Left unclamped deliberately: going over a self-imposed budget is the one // thing it exists to tell you, and the bar clamps its own width anyway. private static func tier(tokens: Int, budget: Int, window: DateInterval) -> QuotaTier? { diff --git a/panel/SessionUsage.swift b/panel/SessionUsage.swift index 8cc2184..12e8399 100644 --- a/panel/SessionUsage.swift +++ b/panel/SessionUsage.swift @@ -131,6 +131,9 @@ struct UsageView: View { } else { FooterHint(label: "Scroll", keys: ["↑↓"]) FooterHint(label: "Top/Bottom", keys: ["⌘↑↓"]) + if nav.selectedUsageClient == .pi { + FooterHint(label: "Window", keys: ["W"]) + } } if nav.usagePane != UsagePane.allCases.last { FooterHint(label: UsagePane.allCases.last?.label ?? "Next", keys: ["→"]) @@ -393,23 +396,7 @@ struct UsageView: View { } } case .pi: - if let pi = nav.piQuota { - // Every row names its denominator: these are the user's own - // budgets, not a limit pi will enforce, and a bare percentage - // here would read like Claude's above it. - if let tier = pi.apiToday { - section("API models today") { tierRow(tier, budget: pi.budget.apiDaily) } - } - if let tier = pi.apiThisWeek { - section("API models this week") { tierRow(tier, budget: pi.budget.apiWeekly) } - } - if let tier = pi.localToday { - section("Local models today") { tierRow(tier, budget: pi.budget.localDaily) } - } - if let tier = pi.localThisWeek { - section("Local models this week") { tierRow(tier, budget: pi.budget.localWeekly) } - } - } + if let pi = nav.piQuota { piPage(pi) } } } @@ -437,15 +424,7 @@ struct UsageView: View { .textCase(.uppercase) // The window lives here rather than in the pane label, since // W cycles it. - Text(nav.usageWindow.label) - .font(.caption2.monospacedDigit().weight(.semibold)) - .foregroundStyle(Color.green) - .padding(.horizontal, 5) - .padding(.vertical, 1) - .background( - RoundedRectangle(cornerRadius: 4, style: .continuous) - .fill(Color.green.opacity(0.14)) - ) + windowPill(nav.usageWindow.label) Spacer() Text(metricValue(series.total(for: metric), metric)) .font(.caption.monospacedDigit().weight(.semibold)) @@ -571,7 +550,48 @@ struct UsageView: View { } } - private func tierRow(_ tier: QuotaTier, budget: Int? = nil) -> some View { + // One model per row, one window at a time. Its own page rather than another + // branch of the shared tiers: pi's limits are the user's, so each row names + // the budget it's measured against, and W swaps the window in place. + @ViewBuilder private func piPage(_ pi: PiQuotaSnapshot) -> some View { + let window = nav.piWindow + let models = pi.models(in: window) + HStack { + windowPill(window.label) + .contentShape(Rectangle()) + .onTapGesture { nav.cyclePiWindow() } + Spacer() + } + .padding(.horizontal, 6) + if models.isEmpty { + Text("No pi usage \(window == .today ? "today" : "this week") yet.") + .font(.caption) + .foregroundStyle(.secondary) + .padding(.horizontal, 6) + } + ForEach(models, id: \.key) { model in + let allowance = pi.budget.allowance(isLocal: model.key.isLocal, in: window) + section(model.name) { + tierRow(model.tier, + caption: "of \(TokenFormat.short(allowance)) \(model.key.isLocal ? "local" : "API") budget") + } + } + } + + // The green chip naming the active window, wherever W cycles it. + private func windowPill(_ label: String) -> some View { + Text(label) + .font(.caption2.monospacedDigit().weight(.semibold)) + .foregroundStyle(Color.green) + .padding(.horizontal, 5) + .padding(.vertical, 1) + .background( + RoundedRectangle(cornerRadius: 4, style: .continuous) + .fill(Color.green.opacity(0.14)) + ) + } + + private func tierRow(_ tier: QuotaTier, caption: String? = nil) -> some View { // Show "30% used" or "70% remaining" depending on the toggle. Bar // still represents utilization so the color ramp keeps its meaning. let display = nav.quotaShowRemaining @@ -581,8 +601,8 @@ struct UsageView: View { return VStack(alignment: .leading, spacing: 4) { HStack(alignment: .firstTextBaseline, spacing: 8) { Spacer() - if let budget { - Text("of \(TokenFormat.short(budget)) budget") + if let caption { + Text(caption) .font(.caption2) .foregroundStyle(.tertiary) } diff --git a/panel/UsageHistory.swift b/panel/UsageHistory.swift index 940251c..701fc06 100644 --- a/panel/UsageHistory.swift +++ b/panel/UsageHistory.swift @@ -117,13 +117,13 @@ struct UsageSeries: Equatable { } } -// Local/API token split over an arbitrary span. The budget's windows are -// calendar days and weeks, which no bucket grid lines up with and which run -// wider than any UsageWindow, so it sums entries rather than buckets. -struct UsageTotals: Equatable { - var localTokens = 0 - var apiTokens = 0 - var turns = 0 +// Keyed on provider and lane as well as the name: pi can reach one model name +// through two providers, and the same model can be free locally and billed +// through an API. +struct UsageModelKey: Hashable { + let provider: String? + let model: String + let isLocal: Bool } // MARK: - Parsed entry @@ -143,6 +143,10 @@ struct UsageEntry: Equatable { // Whether the turn ran on a model that costs nothing. Only pi has one, so // Claude and Codex entries are always billed. let isLocal: Bool + // Only pi's entries carry these; the Usage tab breaks its budget down by + // model. nil for Claude and Codex. + let model: String? + let provider: String? } extension UsageEntry { @@ -320,26 +324,26 @@ final class UsageHistoryStore { return accumulator.series() } - // Sum one source's cached entries over [from, to), split by local and API. - // Deduplicated by identity the same way bucketing is, and pure computation - // over what refresh already cached — the caller is responsible for having - // retained a span this wide. - func totals(source: Source, from: Date, to: Date) -> UsageTotals { + // Tokens per model over [from, to). Sums entries rather than buckets: the pi + // budget's windows are calendar days and weeks, which no bucket grid lines up + // with and which run wider than any UsageWindow. Deduplicated by identity the + // same way bucketing is, and pure computation over what refresh already + // cached, so the caller must have retained a span this wide. Entries with no + // model (Claude, Codex) are skipped. + func totals(source: Source, from: Date, to: Date) -> [UsageModelKey: Int] { lock.lock() let all = (cache[source] ?? [:]).values.flatMap { $0.entries } lock.unlock() var seen: Set = [] - var result = UsageTotals() + var result: [UsageModelKey: Int] = [:] for entry in all { if let identity = entry.identity, !seen.insert(identity).inserted { continue } - guard entry.when >= from, entry.when < to else { continue } - result.turns += 1 - if entry.isLocal { - result.localTokens += entry.tokensExcludingCacheReads - } else { - result.apiTokens += entry.tokensExcludingCacheReads - } + guard entry.when >= from, entry.when < to, + let model = entry.model, entry.tokensExcludingCacheReads > 0 + else { continue } + let key = UsageModelKey(provider: entry.provider, model: model, isLocal: entry.isLocal) + result[key, default: 0] += entry.tokensExcludingCacheReads } return result } @@ -453,7 +457,9 @@ final class UsageHistoryStore { outputTokens: output + reasoning, cacheReadTokens: cacheRead, cacheCreateTokens: cacheWrite, - isLocal: cost <= 0 + isLocal: cost <= 0, + model: message["model"] as? String ?? "unknown", + provider: message["provider"] as? String ) } @@ -472,7 +478,9 @@ final class UsageHistoryStore { outputTokens: usage["output_tokens"] as? Int ?? 0, cacheReadTokens: usage["cache_read_input_tokens"] as? Int ?? 0, cacheCreateTokens: usage["cache_creation_input_tokens"] as? Int ?? 0, - isLocal: false + isLocal: false, + model: nil, + provider: nil ) } @@ -502,7 +510,9 @@ final class UsageHistoryStore { outputTokens: output + reasoning, cacheReadTokens: cached, cacheCreateTokens: 0, - isLocal: false + isLocal: false, + model: nil, + provider: nil ) } diff --git a/panel/WidgetQuota.swift b/panel/WidgetQuota.swift index 73ba037..5aa4ee4 100644 --- a/panel/WidgetQuota.swift +++ b/panel/WidgetQuota.swift @@ -46,7 +46,7 @@ struct WidgetQuota: Equatable { case .antigravity: return "Inner ring: model closest to its limit · Outer ring: monthly prompt credits" case .pi: - return "Inner ring: today's budget · Outer ring: this week's budget" + return "Inner ring: model closest to today's budget · Outer ring: model closest to this week's" default: let inner = "Inner ring: \(Self.ringPhrase(shortLabel, short))" guard long != nil else { return inner } @@ -106,14 +106,16 @@ struct WidgetQuota: Equatable { long: creditsTier(antigravity?.promptCredits), shortLabel: "now", longLabel: "mo") case .pi: - // Both rings read against the user's own budget. API takes the ring - // whenever it has usage: local tokens are free, so an API overrun is - // the one worth a glance. Labels are named rather than measured from - // the window — a DST day is 23 hours long and "23h" in a 35pt legend - // slot would be noise. + // As with Antigravity, the model closest to its budget takes each + // ring, since it's the one about to run out. Labels are named rather + // than measured from the window: a DST day is 23 hours long, and + // "23h" in a 35pt legend slot would be noise. + let closest: ([PiModelUsage]?) -> QuotaTier? = { models in + models?.max { $0.tier.utilization < $1.tier.utilization }?.tier + } return WidgetQuota(client: .pi, - short: pi?.apiToday ?? pi?.localToday, - long: pi?.apiThisWeek ?? pi?.localThisWeek, + short: closest(pi?.today), + long: closest(pi?.thisWeek), shortLabel: "1d", longLabel: "7d") case nil: return .empty