diff --git a/Tests/StackNudgePanelCoreTests/ClaudeCliQuotaProbeTests.swift b/Tests/StackNudgePanelCoreTests/ClaudeCliQuotaProbeTests.swift new file mode 100644 index 0000000..3c9852c --- /dev/null +++ b/Tests/StackNudgePanelCoreTests/ClaudeCliQuotaProbeTests.swift @@ -0,0 +1,224 @@ +import XCTest + +@testable import StackNudgePanelCore + +final class ClaudeCliQuotaProbeTests: XCTestCase { + + // MARK: - parseTierLine + + func testParseTierLineWithResetsSuffix() { + let line = "Current session: 2% used · resets Jun 30 at 6:50pm (Europe/London)" + let parsed = ClaudeCliQuotaProbe.parseTierLine(line) + XCTAssertEqual(parsed?.name, "session") + XCTAssertEqual(parsed?.tier.utilization, 2) + XCTAssertNotNil(parsed?.tier.resetsAt) + } + + func testParseTierLineWithoutResetsSuffix() { + // 0% lines omit the "· resets" suffix. + let line = "Current week (Sonnet only): 0% used" + let parsed = ClaudeCliQuotaProbe.parseTierLine(line) + XCTAssertEqual(parsed?.name, "week (Sonnet only)") + XCTAssertEqual(parsed?.tier.utilization, 0) + XCTAssertNil(parsed?.tier.resetsAt) + } + + func testParseTierLineRejectsMalformed() { + XCTAssertNil(ClaudeCliQuotaProbe.parseTierLine("not a tier line")) + XCTAssertNil(ClaudeCliQuotaProbe.parseTierLine("Current session: pct used")) + XCTAssertNil(ClaudeCliQuotaProbe.parseTierLine("Current: 5% used")) + } + + // MARK: - parseResultText + + func testParseResultTextFullOutput() { + let text = """ + You are currently using your subscription to power your Claude Code usage + + Current session: 2% used · resets Jun 30 at 6:50pm (Europe/London) + Current week (all models): 23% used · resets Jul 4 at 3am (Europe/London) + Current week (Sonnet only): 0% used + Current week (Opus only): 12% used · resets Jul 4 at 3am (Europe/London) + + What's contributing to your limits usage? + Approximate, based on local sessions on this machine — does not include other devices or claude.ai. + + Last 24h · 1346 requests · 12 sessions + 93% of your usage was at >150k context + """ + switch ClaudeCliQuotaProbe.parseResultText(text) { + case .ok(let snap): + XCTAssertEqual(snap.fiveHour?.utilization, 2) + XCTAssertEqual(snap.sevenDay?.utilization, 23) + XCTAssertEqual(snap.sevenDaySonnet?.utilization, 0) + XCTAssertEqual(snap.sevenDayOpus?.utilization, 12) + XCTAssertNil(snap.planType, "planType is filled in by fetch(), not parseResultText") + default: + XCTFail("Expected .ok") + } + } + + func testParseResultTextSoftFailWhenBucketsAbsent() { + // Rate-limited / cold-cache shape — the "What's contributing" section + // still renders but the server-fetched Current lines are missing. + let text = """ + You are currently using your subscription to power your Claude Code usage + + What's contributing to your limits usage? + Approximate, based on local sessions on this machine. + + Last 24h · 1346 requests · 12 sessions + 93% of your usage was at >150k context + """ + switch ClaudeCliQuotaProbe.parseResultText(text) { + case .softFail: break + default: XCTFail("Expected .softFail") + } + } + + func testParseResultTextSkipsMalformedLines() { + // Only the well-formed Sonnet line should land in the snapshot. + let text = """ + Current session: not a number used + Current week (Sonnet only): 7% used + Current: orphaned + """ + switch ClaudeCliQuotaProbe.parseResultText(text) { + case .ok(let snap): + XCTAssertEqual(snap.sevenDaySonnet?.utilization, 7) + XCTAssertNil(snap.fiveHour) + XCTAssertNil(snap.sevenDay) + XCTAssertNil(snap.sevenDayOpus) + default: + XCTFail("Expected .ok") + } + } + + func testParseResultTextIgnoresUnknownBucketNames() { + // A new bucket label Anthropic adds in a future release shouldn't crash + // us — it just gets dropped. Known lines around it must still parse. + let text = """ + Current month (experimental): 50% used · resets Aug 1 at 12am (UTC) + Current session: 8% used · resets Jul 1 at 9am (UTC) + """ + switch ClaudeCliQuotaProbe.parseResultText(text) { + case .ok(let snap): + XCTAssertEqual(snap.fiveHour?.utilization, 8) + XCTAssertNil(snap.sevenDay) + default: + XCTFail("Expected .ok") + } + } + + // MARK: - parseEnvelope + + func testParseEnvelopeUnwrapsResultField() { + let raw = #"{"type":"result","result":"Current session: 5% used\n"}"# + switch ClaudeCliQuotaProbe.parseEnvelope(raw) { + case .ok(let snap): + XCTAssertEqual(snap.fiveHour?.utilization, 5) + default: + XCTFail("Expected .ok") + } + } + + func testParseEnvelopeRejectsNonJson() { + switch ClaudeCliQuotaProbe.parseEnvelope("not json") { + case .hardFail: break + default: XCTFail("Expected .hardFail") + } + } + + func testParseEnvelopeRejectsEmpty() { + switch ClaudeCliQuotaProbe.parseEnvelope("") { + case .hardFail: break + default: XCTFail("Expected .hardFail") + } + switch ClaudeCliQuotaProbe.parseEnvelope(nil) { + case .hardFail: break + default: XCTFail("Expected .hardFail") + } + } + + func testParseEnvelopeRejectsMissingResultField() { + let raw = #"{"type":"result","other":"fields"}"# + switch ClaudeCliQuotaProbe.parseEnvelope(raw) { + case .hardFail: break + default: XCTFail("Expected .hardFail") + } + } + + // MARK: - parseResetsAt + + func testParseResetsAtWithTimezone() { + // 6:50pm in Europe/London on a date in the current year. + let parsed = ClaudeCliQuotaProbe.parseResetsAt("Jun 30 at 6:50pm (Europe/London)") + XCTAssertNotNil(parsed) + var cal = Calendar(identifier: .gregorian) + cal.timeZone = TimeZone(identifier: "Europe/London")! + let comps = cal.dateComponents([.month, .day, .hour, .minute], from: parsed!) + XCTAssertEqual(comps.month, 6) + XCTAssertEqual(comps.day, 30) + XCTAssertEqual(comps.hour, 18) + XCTAssertEqual(comps.minute, 50) + } + + func testParseResetsAtBareHour() { + // "3am" — no minutes — uses the alternate format. + let parsed = ClaudeCliQuotaProbe.parseResetsAt("Jul 4 at 3am (Europe/London)") + XCTAssertNotNil(parsed) + var cal = Calendar(identifier: .gregorian) + cal.timeZone = TimeZone(identifier: "Europe/London")! + let comps = cal.dateComponents([.month, .day, .hour, .minute], from: parsed!) + XCTAssertEqual(comps.month, 7) + XCTAssertEqual(comps.day, 4) + XCTAssertEqual(comps.hour, 3) + XCTAssertEqual(comps.minute, 0) + } + + func testParseResetsAtMalformed() { + XCTAssertNil(ClaudeCliQuotaProbe.parseResetsAt("nonsense")) + XCTAssertNil(ClaudeCliQuotaProbe.parseResetsAt("")) + } + + // MARK: - Session cleanup + + func testExtractSessionId() { + let raw = #"{"type":"result","session_id":"abc-123","result":"x"}"# + XCTAssertEqual(ClaudeCliQuotaProbe.extractSessionId(raw), "abc-123") + } + + func testExtractSessionIdMissing() { + XCTAssertNil(ClaudeCliQuotaProbe.extractSessionId(#"{"type":"result"}"#)) + XCTAssertNil(ClaudeCliQuotaProbe.extractSessionId("garbage")) + XCTAssertNil(ClaudeCliQuotaProbe.extractSessionId(nil)) + } + + func testProbeSessionsDirEncodesCwd() { + // "/Users/hiskud/.stack-nudge" → "-Users-hiskud--stack-nudge". + // Claude replaces BOTH '/' and '.' with '-' (double dash for the + // dot-prefix). The probe directory must match that encoding exactly, + // or the cleanup-by-path will silently miss. + let dir = ClaudeCliQuotaProbe.probeSessionsDir + let encoded = ClaudeCliQuotaProbe.probeCwd + .replacingOccurrences(of: "/", with: "-") + .replacingOccurrences(of: ".", with: "-") + XCTAssertTrue(dir.hasSuffix("/.claude/projects/\(encoded)")) + XCTAssertTrue(encoded.contains("--"), + "encoding must double-dash the dot-prefix on ~/.stack-nudge") + } + + func testRemoveSessionFileRejectsPathTraversal() { + // Defence in depth — a bogus session_id with slashes / .. must NOT + // resolve to anything outside probeSessionsDir. Easiest assertion is + // that the canary file we plant survives the call. + let tmp = FileManager.default.temporaryDirectory + .appendingPathComponent("stack-nudge-canary-\(UUID().uuidString)") + try? "canary".write(to: tmp, atomically: true, encoding: .utf8) + defer { try? FileManager.default.removeItem(at: tmp) } + + ClaudeCliQuotaProbe.removeSessionFile("../../../../../../\(tmp.lastPathComponent)") + XCTAssertTrue(FileManager.default.fileExists(atPath: tmp.path), + "removeSessionFile must reject anything outside its directory") + } +} diff --git a/panel/ClaudeCliQuotaProbe.swift b/panel/ClaudeCliQuotaProbe.swift new file mode 100644 index 0000000..3994d24 --- /dev/null +++ b/panel/ClaudeCliQuotaProbe.swift @@ -0,0 +1,228 @@ +import Foundation + +// Reads Claude Code's quota by shelling out to `claude --print --output-format +// json /usage` and parsing the human-readable text the slash command renders. +// /usage is a client-side intercept (zero model cost, num_turns 0, +// duration_api_ms 0) that hits the same /api/oauth/usage endpoint our legacy +// QuotaProbe does — but it runs inside Claude Code, which has its own keychain +// ACL grant. Our process never touches the keychain, so the periodic +// "stack-nudge wants to access Claude Code-credentials" prompt goes away. +// +// Failure shapes: +// hardFail — `claude` not on PATH, spawn/timeout failure, or unparseable +// JSON envelope. PanelController falls back to the legacy probe. +// softFail — CLI ran fine but the server-side bucket lines are absent +// (rate-limited cold cache). We back off 60s locally without +// clearing the existing snapshot; next tick should populate. +// ok — at least one Current bucket line parsed. +final class ClaudeCliQuotaProbe { + + // Main-queue only (mirrors QuotaProbe's threading model). + private(set) var lastProbeFailed = false + private var retryAfterUntil: Date? + private var lastSubscriptionType: String? + private var subscriptionFetched = false + + var isRateLimited: Bool { + guard let until = retryAfterUntil else { return false } + return until > Date() + } + + // CLI shell-out runs off-main; completion fires on main. + private let probeQueue = DispatchQueue(label: "stack-nudge.claude-cli-quota") + + func fetch(completion: @escaping (QuotaSnapshot?) -> Void) { + if isRateLimited { + completion(nil) + return + } + guard let path = ProcessOutput.claude() else { + lastProbeFailed = true + completion(nil) + return + } + let needsSubscriptionFetch = !subscriptionFetched + let priorPlan = lastSubscriptionType + + probeQueue.async { [weak self] in + var fetchedPlan: String? = priorPlan + if needsSubscriptionFetch { + if let json = ProcessOutput.read( + path, ["auth", "status", "--json"], timeout: 5, cwd: Self.probeCwd), + let data = json.data(using: .utf8), + let obj = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let plan = obj["subscriptionType"] as? String { + fetchedPlan = plan + } + } + let raw = ProcessOutput.read( + path, ["--print", "--output-format", "json", "/usage"], + timeout: 10, cwd: Self.probeCwd) + let result = Self.parseEnvelope(raw) + // Every `claude --print` spawns a new session rollout under + // ~/.claude/projects//.jsonl. At a 60s poll + // cadence that's ~1.4k files/day — clutters `claude --resume` + // and burns inodes for zero benefit (num_turns = 0). Pinning + // cwd above means we know the exact directory; this scrubs the + // single file the probe just produced. + if let sid = Self.extractSessionId(raw) { + Self.removeSessionFile(sid) + } + + DispatchQueue.main.async { + guard let self else { return } + if needsSubscriptionFetch { self.subscriptionFetched = true } + self.lastSubscriptionType = fetchedPlan + + switch result { + case .ok(let snap): + self.lastProbeFailed = false + self.retryAfterUntil = nil + completion(QuotaSnapshot( + fiveHour: snap.fiveHour, + sevenDay: snap.sevenDay, + sevenDayOpus: snap.sevenDayOpus, + sevenDaySonnet: snap.sevenDaySonnet, + planType: fetchedPlan)) + case .softFail: + self.lastProbeFailed = false + self.retryAfterUntil = Date().addingTimeInterval(60) + completion(nil) + case .hardFail: + self.lastProbeFailed = true + completion(nil) + } + } + } + } + + // MARK: - Session cleanup + + // Pinned cwd for `claude --print` so its session-rollout file always + // lands in a known directory we can scrub. ~/.stack-nudge always exists + // by the time the probe runs (Bootstrap creates it on first launch). + static let probeCwd = "\(NSHomeDirectory())/.stack-nudge" + + // Claude derives the projects-subdirectory name by replacing every "/" + // AND every "." in the cwd with "-". So "/Users/me/.stack-nudge" becomes + // "-Users-me--stack-nudge" (double dash for the dot-prefix). + static var probeSessionsDir: String { + let encoded = probeCwd + .replacingOccurrences(of: "/", with: "-") + .replacingOccurrences(of: ".", with: "-") + return "\(NSHomeDirectory())/.claude/projects/\(encoded)" + } + + static func extractSessionId(_ raw: String?) -> String? { + guard let raw, + let data = raw.data(using: .utf8), + let obj = try? JSONSerialization.jsonObject(with: data) as? [String: Any] + else { return nil } + return obj["session_id"] as? String + } + + static func removeSessionFile(_ sessionId: String) { + // Guard against a malformed session_id that could escape the + // intended directory — only accept the UUID shape Claude actually + // emits (lower-hex + hyphens, no separators). + guard sessionId.range(of: "^[0-9a-f-]+$", options: .regularExpression) != nil + else { return } + let path = "\(probeSessionsDir)/\(sessionId).jsonl" + try? FileManager.default.removeItem(atPath: path) + } + + // MARK: - Parsing + + enum ParseResult { + case ok(QuotaSnapshot) // planType is filled in by the caller + case softFail // bucket lines absent — CLI is rate-limited + case hardFail // envelope missing / unparseable + } + + static func parseEnvelope(_ raw: String?) -> ParseResult { + guard let raw, !raw.isEmpty, + let data = raw.data(using: .utf8), + let obj = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let text = obj["result"] as? String else { + return .hardFail + } + return parseResultText(text) + } + + static func parseResultText(_ text: String) -> ParseResult { + var fiveHour: QuotaTier? + var sevenDay: QuotaTier? + var opus: QuotaTier? + var sonnet: QuotaTier? + var foundAny = false + + for line in text.split(separator: "\n") { + let s = String(line) + guard s.hasPrefix("Current "), + let parsed = parseTierLine(s) else { continue } + foundAny = true + switch parsed.name { + case "session": fiveHour = parsed.tier + case "week (all models)": sevenDay = parsed.tier + case "week (Opus only)": opus = parsed.tier + case "week (Sonnet only)": sonnet = parsed.tier + default: continue + } + } + if !foundAny { return .softFail } + return .ok(QuotaSnapshot( + fiveHour: fiveHour, + sevenDay: sevenDay, + sevenDayOpus: opus, + sevenDaySonnet: sonnet, + planType: nil)) + } + + // "Current week (all models): 23% used · resets Jul 4 at 3am (Europe/London)" + // "Current week (Sonnet only): 0% used" ← no resets suffix on 0% + private static let lineRegex = try! NSRegularExpression( + pattern: #"^Current (.+?): (\d+)% used(?: · resets (.+))?$"#) + + static func parseTierLine(_ line: String) -> (name: String, tier: QuotaTier)? { + let ns = line as NSString + let range = NSRange(location: 0, length: ns.length) + guard let m = lineRegex.firstMatch(in: line, range: range) else { return nil } + let name = ns.substring(with: m.range(at: 1)) + guard let pct = Double(ns.substring(with: m.range(at: 2))) else { return nil } + let resetsAt: Date? = { + let r = m.range(at: 3) + guard r.location != NSNotFound else { return nil } + return parseResetsAt(ns.substring(with: r)) + }() + return (name, QuotaTier(utilization: pct, resetsAt: resetsAt)) + } + + // Best-effort: "Jun 30 at 6:50pm (Europe/London)" → Date. + // The CLI omits the year, so we splice in the current one. If the deadline + // is "Jan 2 at 1am" rendered on Dec 31, the parsed date will be ~12 months + // in the past — accept the rare wraparound, the UI just shows nil when + // resetsAt is in the past (RelativeTime treats negative deltas gracefully). + static func parseResetsAt(_ raw: String) -> Date? { + let trimmed = raw.trimmingCharacters(in: .whitespaces) + var dateTimeStr = trimmed + var tz: TimeZone? + if trimmed.hasSuffix(")"), let openIdx = trimmed.lastIndex(of: "(") { + let afterOpen = trimmed.index(after: openIdx) + let beforeClose = trimmed.index(before: trimmed.endIndex) + let tzString = String(trimmed[afterOpen.. String? { + let task = Process() + task.executableURL = URL(fileURLWithPath: path) + task.arguments = args + if let cwd { task.currentDirectoryURL = URL(fileURLWithPath: cwd) } + let outPipe = Pipe() + task.standardOutput = outPipe + task.standardError = Pipe() + + let group = DispatchGroup() + group.enter() + task.terminationHandler = { _ in group.leave() } + + do { try task.run() } catch { + // Balance the enter() — terminationHandler never fires if run() throws. + group.leave() + return nil + } + + if group.wait(timeout: .now() + timeout) == .timedOut { + task.terminate() + if group.wait(timeout: .now() + 1) == .timedOut, task.isRunning { + kill(task.processIdentifier, SIGKILL) + _ = group.wait(timeout: .now() + 1) + } + return nil + } + let data = outPipe.fileHandleForReading.readDataToEndOfFile() + return String(data: data, encoding: .utf8) ?? "" + } + // Resolve the `gh` CLI from common install locations. A launchd-spawned app // has a minimal PATH, so we probe paths directly rather than relying on env. // nil ⇒ not installed (callers no-op). Used for release checks/downloads. @@ -26,4 +63,16 @@ enum ProcessOutput { ["/opt/homebrew/bin/gh", "/usr/local/bin/gh", "/usr/bin/gh"] .first { FileManager.default.isExecutableFile(atPath: $0) } } + + // Resolve the `claude` CLI. Same minimal-PATH rationale as gh(); the + // ~/.claude/local fallback covers users who installed via the curl-bash + // installer rather than Homebrew. + static func claude() -> String? { + let home = NSHomeDirectory() + return [ + "/opt/homebrew/bin/claude", + "/usr/local/bin/claude", + "\(home)/.claude/local/claude", + ].first { FileManager.default.isExecutableFile(atPath: $0) } + } } diff --git a/panel/Settings.swift b/panel/Settings.swift index c6e0686..85c4111 100644 --- a/panel/Settings.swift +++ b/panel/Settings.swift @@ -84,13 +84,8 @@ struct SettingsView: View { row(.pollFrequency, label: "Poll frequency", kind: .cycle, value: "\(nav.quotaPollMinutes) min", enabled: nav.quotaTrackingEnabled) row(.contextAlert, label: "Context alert at", kind: .cycle, value: contextAlertLabel) row(.showRemaining, label: "Show remaining", kind: .toggle, value: nav.quotaShowRemaining ? "On" : "Off", enabled: nav.quotaTrackingEnabled) - if nav.usingPlaintextCredentials { - Text("⚠︎ Reading the Claude token from ~/.claude/.credentials.json (plaintext) — any process running as you can read it. Delete that file to fall back to the Keychain.") - .font(.caption) - .foregroundStyle(.orange) - .padding(.horizontal, 14) - .padding(.top, 2) - .fixedSize(horizontal: false, vertical: true) + if nav.quotaTrackingEnabled { + probeSourceLabel } } @@ -317,6 +312,28 @@ struct SettingsView: View { nav.contextAlertThresholdK == 0 ? "Off" : "\(nav.contextAlertThresholdK)K" } + // Only surfaced when there's something the user needs to know: the + // plaintext-file security tradeoff, or the keychain-rotation prompts. + // The CLI probe path (the default) is silent — it's the happy state. + @ViewBuilder private var probeSourceLabel: some View { + if nav.usingPlaintextCredentials { + label("⚠︎ Reading the Claude token from ~/.claude/.credentials.json (plaintext) — any process running as you can read it.", + color: .orange) + } else if !nav.usingClaudeCliProbe { + label("Reading via macOS Keychain — periodic password prompts are expected when Claude Code rotates its token.", + color: .secondary) + } + } + + private func label(_ text: String, color: Color) -> some View { + Text(text) + .font(.caption) + .foregroundStyle(color) + .padding(.horizontal, 14) + .padding(.top, 2) + .fixedSize(horizontal: false, vertical: true) + } + private var checkForUpdatesStatus: String { switch nav.updateCheckStatus { case .idle: return ""