diff --git a/Sources/OpenUsage/Providers/Codex/CodexHistoryRefresh.swift b/Sources/OpenUsage/Providers/Codex/CodexHistoryRefresh.swift new file mode 100644 index 000000000..444725ce0 --- /dev/null +++ b/Sources/OpenUsage/Providers/Codex/CodexHistoryRefresh.swift @@ -0,0 +1,33 @@ +import Foundation + +/// A single local scan can outlive a quota refresh. Never start overlapping scans, and consume the +/// completed result on the next refresh if it misses this one's short wait budget. Unlike a task +/// group, this bounded wait does not wait for a scanner that ignores cancellation to unwind. +@MainActor +final class CodexHistoryRefresh { + private var task: Task? + private var completed: Value? + + deinit { task?.cancel() } + + func value(wait: Duration, operation: @escaping @MainActor () async -> Value) async -> Value? { + guard !Task.isCancelled else { return nil } + if task == nil { + task = Task { [weak self] in + let value = await operation() + guard !Task.isCancelled else { return } + self?.completed = value + } + } + let clock = ContinuousClock() + let deadline = clock.now.advanced(by: wait) + while completed == nil && clock.now < deadline && !Task.isCancelled { + // Only active while waiting for a scan, never an idle application timer. + try? await Task.sleep(for: min(.milliseconds(25), clock.now.duration(to: deadline))) + } + guard !Task.isCancelled, let result = completed else { return nil } + completed = nil + task = nil + return result + } +} diff --git a/Sources/OpenUsage/Providers/Codex/CodexProvider.swift b/Sources/OpenUsage/Providers/Codex/CodexProvider.swift index 5497dd662..d58e0ad44 100644 --- a/Sources/OpenUsage/Providers/Codex/CodexProvider.swift +++ b/Sources/OpenUsage/Providers/Codex/CodexProvider.swift @@ -12,6 +12,9 @@ final class CodexProvider: ProviderRuntime { ] ) + private let localHistory = CodexHistoryRefresh() + let localHistoryWait: Duration + let authStore: CodexAuthStore let usageClient: CodexUsageClient let logUsageScanner: CodexLogUsageScanner @@ -21,6 +24,7 @@ final class CodexProvider: ProviderRuntime { let fallbackModel: @MainActor () -> String? init( + localHistoryWait: Duration = .seconds(2), authStore: CodexAuthStore = CodexAuthStore(), usageClient: CodexUsageClient = CodexUsageClient(), logUsageScanner: CodexLogUsageScanner = CodexLogUsageScanner(), @@ -29,6 +33,7 @@ final class CodexProvider: ProviderRuntime { pricing: @escaping @Sendable () async -> ModelPricing = { await ModelPricingStore.shared.current() }, fallbackModel: @escaping @MainActor () -> String? = { CodexFallbackModelSetting.current() } ) { + self.localHistoryWait = localHistoryWait self.authStore = authStore self.usageClient = usageClient self.logUsageScanner = logUsageScanner @@ -142,12 +147,35 @@ final class CodexProvider: ProviderRuntime { ) var mapped = try CodexUsageMapper.mapUsageResponse(response, resetCredits: resetCredits, now: now()) - // Local spend tiles, scanned natively from the Codex CLI's session rollouts and priced through - // the shared pricing store, merged with Codex usage that happened inside pi or OpenCode. Those - // agents attribute their underlying Codex OAuth traffic back to this card. + let history = await localHistory.value(wait: localHistoryWait) { [self] in + await scanLocalHistory() + } + if let history { mapped.lines += history.lines } + let warning = history == nil ? "Local token history is still updating." : nil + if warning != nil { + AppLog.warn(LogTag.plugin("codex"), "local history scan deferred; publishing live quota") + } + MetricLine.appendNoDataIfNeeded(&mapped.lines) + return ProviderSnapshot.make( + provider: provider, + plan: mapped.plan, + lines: mapped.lines, + refreshedAt: now(), + usageHistory: history?.usageHistory, + warning: warning + ) + } + + private struct CodexLocalHistory: Sendable { + var lines: [MetricLine] + var usageHistory: ProviderUsageHistory? + } + + private func scanLocalHistory() async -> CodexLocalHistory { + var lines: [MetricLine] = [] let pricing = await pricing() // Three independent local sources: reading rollout files, pi's JSONL, and OpenCode's SQLite - // concurrently keeps the slowest one — not their sum — on the refresh's critical path. + // concurrently keeps the slowest one — not their sum — on the background scan's path. let selectedFallbackModel = fallbackModel() async let native = logUsageScanner.scan( now: now(), pricing: pricing, fallbackModel: selectedFallbackModel @@ -170,26 +198,19 @@ final class CodexProvider: ProviderRuntime { fallbackPricingModelsByDay: scan.fallbackPricingModelsByDay ) SpendTileMapper.appendTokenUsage( - scan.series, to: &mapped.lines, now: now(), + scan.series, to: &lines, now: now(), unknownModelsByDay: scan.unknownModelsByDay, modelUsage: scan.modelUsage, modelSourceNote: baseNote, fallbackPricingModelsByDay: scan.fallbackPricingModelsByDay ) SpendTileMapper.appendUsageTrend( - scan.series, to: &mapped.lines, now: now(), note: baseNote, + scan.series, to: &lines, now: now(), note: baseNote, fallbackPricingModelsByDay: scan.fallbackPricingModelsByDay ) } - MetricLine.appendNoDataIfNeeded(&mapped.lines) - return ProviderSnapshot.make( - provider: provider, - plan: mapped.plan, - lines: mapped.lines, - refreshedAt: now(), - usageHistory: usageHistory - ) + return CodexLocalHistory(lines: lines, usageHistory: usageHistory) } private static func localUsageSourceNote(hasPi: Bool, hasOpenCode: Bool) -> String { diff --git a/Tests/OpenUsageTests/CodexHistoryRefreshTests.swift b/Tests/OpenUsageTests/CodexHistoryRefreshTests.swift new file mode 100644 index 000000000..0c7730d2b --- /dev/null +++ b/Tests/OpenUsageTests/CodexHistoryRefreshTests.swift @@ -0,0 +1,48 @@ +import XCTest +@testable import OpenUsage + +@MainActor +final class CodexHistoryRefreshTests: XCTestCase { + func testSlowScanDoesNotBlockOrRestartAndCompletedResultIsCollected() async { + let refresh = CodexHistoryRefresh() + var continuation: CheckedContinuation? + var starts = 0 + let operation: @MainActor () async -> Int = { + starts += 1 + return await withCheckedContinuation { continuation = $0 } + } + let first = await refresh.value(wait: .zero, operation: operation) + XCTAssertNil(first) + while continuation == nil { await Task.yield() } + let second = await refresh.value(wait: .zero, operation: operation) + XCTAssertNil(second) + XCTAssertEqual(starts, 1) + continuation?.resume(returning: 42) + let completed = await refresh.value(wait: .seconds(1), operation: operation) + XCTAssertEqual(completed, 42) + XCTAssertEqual(starts, 1) + let next = await refresh.value(wait: .seconds(1)) { 43 } + XCTAssertEqual(next, 43) + } + + func testCancelledWaitDoesNotDiscardInFlightScan() async { + let refresh = CodexHistoryRefresh() + var release: CheckedContinuation? + let waiter = Task { await refresh.value(wait: .seconds(30)) { + await withCheckedContinuation { release = $0 } + } } + while release == nil { await Task.yield() } + waiter.cancel() + let cancelled = await waiter.value + XCTAssertNil(cancelled) + release?.resume(returning: 9) + let recovered = await refresh.value(wait: .seconds(1)) { XCTFail("Duplicate scan"); return 0 } + XCTAssertEqual(recovered, 9) + } + + func testFastScanIsIncludedInSameRefresh() async { + let refresh = CodexHistoryRefresh() + let value = await refresh.value(wait: .seconds(1)) { 7 } + XCTAssertEqual(value, 7) + } +} diff --git a/Tests/OpenUsageTests/CodexProviderTests.swift b/Tests/OpenUsageTests/CodexProviderTests.swift index ae175bb5e..fd612e011 100644 --- a/Tests/OpenUsageTests/CodexProviderTests.swift +++ b/Tests/OpenUsageTests/CodexProviderTests.swift @@ -490,6 +490,37 @@ final class CodexUsageMapperTests: XCTestCase { @MainActor final class CodexProviderTests: XCTestCase { + func testLiveQuotaReturnsWhileLocalHistoryIsStillRunning() async throws { + let home = try CodexLogFixture.makeHome(files: [:]) + let provider = CodexProvider( + localHistoryWait: .zero, + authStore: CodexAuthStore( + environment: FakeEnvironment(["CODEX_HOME": "/tmp/codex-fixture"]), + files: FakeFiles(["/tmp/codex-fixture/auth.json": #"{"tokens":{"access_token":"fixture"}}"#]), + keychain: FakeKeychain() + ), + usageClient: CodexUsageClient(http: FakeHTTPClient(response: HTTPResponse( + statusCode: 200, headers: [:], + body: Data(#"{"rate_limit":{"secondary_window":{"used_percent":58,"limit_window_seconds":604800,"reset_after_seconds":3600}}}"#.utf8) + ))), + logUsageScanner: CodexLogFixture.scanner(home: home), + now: { Date(timeIntervalSince1970: 4_075_747_200) }, + pricing: { + try? await Task.sleep(for: .milliseconds(100)) + return ModelPricing(supplement: PricingSupplement(), primary: PricingCatalog(entries: [:]), secondary: PricingCatalog(entries: [:])) + } + ) + let snapshot = await provider.refresh() + XCTAssertNil(snapshot.errorCategory) + XCTAssertNotNil(snapshot.warning) + XCTAssertNil(snapshot.usageHistory) + guard case .progress(_, let used, let limit, _, _, _, _) = snapshot.line(label: "Weekly") else { + return XCTFail("Live weekly quota was lost while local history was pending") + } + XCTAssertEqual(used, 58) + XCTAssertEqual(limit, 100) + } + func testNoUsageDataBadgeIsDroppedWhenLocalLogsHaveSpend() async throws { let now = OpenUsageISO8601.date(from: "2026-02-20T14:30:00.000Z")! // The live usage API returns nothing mappable (empty body -> no metric lines)... diff --git a/docs/providers/codex.md b/docs/providers/codex.md index b6e1ac635..ed93cc7ae 100644 --- a/docs/providers/codex.md +++ b/docs/providers/codex.md @@ -60,3 +60,12 @@ Safeguards, because a claim is irreversible: - Claiming is always a deliberate two-click flow behind the hover popover — nothing is ever claimed automatically. - Each claim targets one explicit credit (re-matched against a fresh credit list at claim time) and carries an idempotency key, so a retry after a network error can never spend a second credit. - If the credit was meanwhile used elsewhere (CLI or web) the popover says it's no longer available and refreshes; if your usage doesn't need a reset, Codex refuses without spending the credit and the popover says so. After a claim resets usage, the remaining Use buttons disable ("nothing to reset") until the popover is reopened. + +### Slow Local History + +Live quota refreshes wait at most two seconds for local token-history processing. If a large local +archive takes longer, quota still updates and the card shows a history-updating notice. The scan +continues in the background; its result is collected by a later refresh. Only one scan runs at a +time. Previously loaded history is retained while waiting. A fresh launch may therefore show quota +before spend/history appears. Network or authentication failures still use the normal stale-data +handling.