From 2fb7bcc39209070feecced3f1a10ad4b43594518 Mon Sep 17 00:00:00 2001 From: thrr87 <193831865+thrr87@users.noreply.github.com> Date: Wed, 29 Jul 2026 16:17:12 +0200 Subject: [PATCH 1/5] perf: reuse timestamp parser during activity restore --- .../CodexLimits/LocalActivityCollector.swift | 3 +- .../LocalActivityPerformanceTests.swift | 71 +++++++++++++++++++ 2 files changed, 73 insertions(+), 1 deletion(-) diff --git a/Sources/CodexLimits/LocalActivityCollector.swift b/Sources/CodexLimits/LocalActivityCollector.swift index e105ef3..65824c9 100644 --- a/Sources/CodexLimits/LocalActivityCollector.swift +++ b/Sources/CodexLimits/LocalActivityCollector.swift @@ -92,6 +92,7 @@ actor LocalActivityCollector { private let installedCLIVersion: (@Sendable () async -> String?)? private let tail = IncrementalRolloutTailSource() private let normalizer = LocalActivityNormalizer() + private let timestampParser = LocalEventTimestampParser() private var files: [String: FileState] = [:] private var restoredFilesByFingerprint: [String: FileState] = [:] private var restoredFingerprintByIdentity: [RolloutFileIdentity: String] = [:] @@ -802,7 +803,7 @@ actor LocalActivityCollector { } private func parseTimestamp(_ value: String) -> Date? { - LocalEventTimestampParser().date(from: value) + timestampParser.date(from: value) } @discardableResult diff --git a/Tests/CodexLimitsTests/LocalActivityPerformanceTests.swift b/Tests/CodexLimitsTests/LocalActivityPerformanceTests.swift index e51390b..4ea5e0c 100644 --- a/Tests/CodexLimitsTests/LocalActivityPerformanceTests.swift +++ b/Tests/CodexLimitsTests/LocalActivityPerformanceTests.swift @@ -4,6 +4,77 @@ import XCTest @testable import CodexLimits final class LocalActivityPerformanceTests: XCTestCase { + func testPersistedFactRestoreStaysResponsive() async throws { + let root = temporaryDirectory() + let rolloutDirectory = root.appendingPathComponent( + "2026/07/28", + isDirectory: true + ) + try FileManager.default.createDirectory( + at: rolloutDirectory, + withIntermediateDirectories: true + ) + let rollout = rolloutDirectory.appendingPathComponent( + "rollout-2026-07-28T10-00-00-benchmark.jsonl" + ) + let recordCount = 20_000 + var fixture = + #"{"timestamp":"2026-07-28T10:00:00.000Z","ordinal":0,"type":"session_meta","payload":{"id":"benchmark","cli_version":"0.145.0"}}"# + + "\n" + fixture.reserveCapacity(recordCount * 180) + for ordinal in 1...recordCount { + fixture += + #"{"timestamp":"2026-07-28T10:00:01.000Z","ordinal":\#(ordinal),"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"total_tokens":\#(ordinal * 100)}}}}"# + + "\n" + } + try Data(fixture.utf8).write(to: rollout) + fixture.removeAll(keepingCapacity: false) + + let stateDirectory = root.appendingPathComponent( + "state", + isDirectory: true + ) + let interval = DateInterval( + start: try XCTUnwrap( + ISO8601DateFormatter().date(from: "2026-07-28T00:00:00Z") + ), + end: try XCTUnwrap( + ISO8601DateFormatter().date(from: "2026-07-29T00:00:00Z") + ) + ) + let first = LocalActivityCollector( + rootDirectory: root, + stateDirectory: stateDirectory + ) + await first.selectPartition("benchmark") + _ = await first.refresh(interval: interval) + + let restarted = LocalActivityCollector( + rootDirectory: root, + stateDirectory: stateDirectory + ) + await restarted.selectPartition("benchmark") + let start = ProcessInfo.processInfo.systemUptime + let restored = await restarted.refresh(interval: interval) + let milliseconds = + (ProcessInfo.processInfo.systemUptime - start) * 1_000 + + print( + String( + format: "PERSISTED_FACT_RESTORE records=%d wall_ms=%.3f", + restored.facts.count, + milliseconds + ) + ) + XCTAssertEqual(restored.bytesRead, 0) + XCTAssertEqual( + restored.facts.filter { $0.key == .token } + .compactMap(\.numericDelta).count, + recordCount - 1 + ) + XCTAssertLessThan(milliseconds, 3_000) + } + func testRepresentativeFixtureMetrics() throws { let directory = temporaryDirectory() try FileManager.default.createDirectory( From aec8ba7147ec7e422cb7e2bc23b180a197c1eca7 Mon Sep 17 00:00:00 2001 From: thrr87 <193831865+thrr87@users.noreply.github.com> Date: Wed, 29 Jul 2026 16:17:12 +0200 Subject: [PATCH 2/5] perf: reuse timestamp parser during activity restore --- .../CodexLimits/LocalActivityCollector.swift | 3 +- .../LocalActivityPerformanceTests.swift | 71 +++++++++++++++++++ 2 files changed, 73 insertions(+), 1 deletion(-) diff --git a/Sources/CodexLimits/LocalActivityCollector.swift b/Sources/CodexLimits/LocalActivityCollector.swift index e105ef3..65824c9 100644 --- a/Sources/CodexLimits/LocalActivityCollector.swift +++ b/Sources/CodexLimits/LocalActivityCollector.swift @@ -92,6 +92,7 @@ actor LocalActivityCollector { private let installedCLIVersion: (@Sendable () async -> String?)? private let tail = IncrementalRolloutTailSource() private let normalizer = LocalActivityNormalizer() + private let timestampParser = LocalEventTimestampParser() private var files: [String: FileState] = [:] private var restoredFilesByFingerprint: [String: FileState] = [:] private var restoredFingerprintByIdentity: [RolloutFileIdentity: String] = [:] @@ -802,7 +803,7 @@ actor LocalActivityCollector { } private func parseTimestamp(_ value: String) -> Date? { - LocalEventTimestampParser().date(from: value) + timestampParser.date(from: value) } @discardableResult diff --git a/Tests/CodexLimitsTests/LocalActivityPerformanceTests.swift b/Tests/CodexLimitsTests/LocalActivityPerformanceTests.swift index e51390b..4ea5e0c 100644 --- a/Tests/CodexLimitsTests/LocalActivityPerformanceTests.swift +++ b/Tests/CodexLimitsTests/LocalActivityPerformanceTests.swift @@ -4,6 +4,77 @@ import XCTest @testable import CodexLimits final class LocalActivityPerformanceTests: XCTestCase { + func testPersistedFactRestoreStaysResponsive() async throws { + let root = temporaryDirectory() + let rolloutDirectory = root.appendingPathComponent( + "2026/07/28", + isDirectory: true + ) + try FileManager.default.createDirectory( + at: rolloutDirectory, + withIntermediateDirectories: true + ) + let rollout = rolloutDirectory.appendingPathComponent( + "rollout-2026-07-28T10-00-00-benchmark.jsonl" + ) + let recordCount = 20_000 + var fixture = + #"{"timestamp":"2026-07-28T10:00:00.000Z","ordinal":0,"type":"session_meta","payload":{"id":"benchmark","cli_version":"0.145.0"}}"# + + "\n" + fixture.reserveCapacity(recordCount * 180) + for ordinal in 1...recordCount { + fixture += + #"{"timestamp":"2026-07-28T10:00:01.000Z","ordinal":\#(ordinal),"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"total_tokens":\#(ordinal * 100)}}}}"# + + "\n" + } + try Data(fixture.utf8).write(to: rollout) + fixture.removeAll(keepingCapacity: false) + + let stateDirectory = root.appendingPathComponent( + "state", + isDirectory: true + ) + let interval = DateInterval( + start: try XCTUnwrap( + ISO8601DateFormatter().date(from: "2026-07-28T00:00:00Z") + ), + end: try XCTUnwrap( + ISO8601DateFormatter().date(from: "2026-07-29T00:00:00Z") + ) + ) + let first = LocalActivityCollector( + rootDirectory: root, + stateDirectory: stateDirectory + ) + await first.selectPartition("benchmark") + _ = await first.refresh(interval: interval) + + let restarted = LocalActivityCollector( + rootDirectory: root, + stateDirectory: stateDirectory + ) + await restarted.selectPartition("benchmark") + let start = ProcessInfo.processInfo.systemUptime + let restored = await restarted.refresh(interval: interval) + let milliseconds = + (ProcessInfo.processInfo.systemUptime - start) * 1_000 + + print( + String( + format: "PERSISTED_FACT_RESTORE records=%d wall_ms=%.3f", + restored.facts.count, + milliseconds + ) + ) + XCTAssertEqual(restored.bytesRead, 0) + XCTAssertEqual( + restored.facts.filter { $0.key == .token } + .compactMap(\.numericDelta).count, + recordCount - 1 + ) + XCTAssertLessThan(milliseconds, 3_000) + } + func testRepresentativeFixtureMetrics() throws { let directory = temporaryDirectory() try FileManager.default.createDirectory( From 2ce19ba3ba0cbb836d05723dbf1c9bb376f5f34b Mon Sep 17 00:00:00 2001 From: thrr87 <193831865+thrr87@users.noreply.github.com> Date: Thu, 30 Jul 2026 00:24:14 +0200 Subject: [PATCH 3/5] perf: bound refresh memory and local ingestion --- .../CodexLimits/ActiveTimeAvailability.swift | 26 +- Sources/CodexLimits/ActivityTimeline.swift | 50 +- Sources/CodexLimits/AnalyticsWorkspace.swift | 70 +- .../CodexLimits/CodexAssistedHistory.swift | 100 +- .../CodexLimits/CodexAssistedInsights.swift | 8 +- Sources/CodexLimits/CodexClient.swift | 62 +- Sources/CodexLimits/ForecastEngine.swift | 51 +- .../CodexLimits/LocalActivityCollector.swift | 998 +++++++++++++++--- .../CodexLimits/LocalActivityFactIndex.swift | 44 +- .../CodexLimits/LocalActivityNormalizer.swift | 35 +- Sources/CodexLimits/LocalTokenActivity.swift | 38 +- Sources/CodexLimits/LocalWorkloadMix.swift | 17 +- Sources/CodexLimits/MenuContentView.swift | 429 +++++--- Sources/CodexLimits/RolloutTailSource.swift | 620 +++++++++-- Sources/CodexLimits/SettingsView.swift | 13 +- Sources/CodexLimits/UsageHistory.swift | 74 +- .../CodexLimits/UsageIntelligenceEngine.swift | 350 +++++- Sources/CodexLimits/UsageModels.swift | 139 +++ Sources/CodexLimits/UsageMonitor.swift | 335 +++++- Sources/CodexLimits/UsagePerToken.swift | 24 +- Sources/CodexLimits/UsageReceipts.swift | 880 ++++++++++++--- .../ActiveTimeAvailabilityTests.swift | 48 +- .../ActivityTimelineTests.swift | 44 + .../CodexAssistedInsightTests.swift | 114 +- Tests/CodexLimitsTests/CodexClientTests.swift | 143 ++- .../ForecastEngineTests.swift | 63 ++ .../LocalActivityCollectorTests.swift | 730 ++++++++++++- .../LocalActivityNormalizerTests.swift | 140 ++- .../LocalActivityPerformanceTests.swift | 209 +++- .../LocalTokenActivityTests.swift | 31 + .../RolloutTailSourceTests.swift | 459 +++++++- .../CodexLimitsTests/UsageHistoryTests.swift | 128 +++ .../UsageIntelligenceEngineTests.swift | 371 ++++++- .../UsageMonitorHistoryTests.swift | 451 ++++++++ .../CodexLimitsTests/UsageReceiptTests.swift | 249 ++++- 35 files changed, 6682 insertions(+), 861 deletions(-) diff --git a/Sources/CodexLimits/ActiveTimeAvailability.swift b/Sources/CodexLimits/ActiveTimeAvailability.swift index 4262d6d..bbcfb28 100644 --- a/Sources/CodexLimits/ActiveTimeAvailability.swift +++ b/Sources/CodexLimits/ActiveTimeAvailability.swift @@ -36,6 +36,10 @@ struct ActiveTimeAvailableEstimate: Equatable, Sendable { struct ActiveTimeAvailabilitySnapshot: Equatable, Sendable { let activeTimeThisWeek: TimeInterval + let maximumConcurrency: Int + let waitingTime: TimeInterval? + let pollingTime: TimeInterval? + let activityBreakdownReason: String? let activeTimeCoverage: CoverageLevel let activeTimeReason: String? let observedInterval: DateInterval? @@ -54,7 +58,8 @@ enum ActiveTimeWeekEvidenceBuilder { usage: [WeeklyUsageEvidence], facts: [LocalActivityFact], projections: [ThreadProjection], - observation: LocalActivityObservation + observation: LocalActivityObservation, + factIndex: LocalActivityFactIndex? = nil ) -> ActiveTimeHistorySelection { guard let currentUsage else { return ActiveTimeHistorySelection( @@ -62,7 +67,7 @@ enum ActiveTimeWeekEvidenceBuilder { unavailableReason: nil ) } - let factIndex = LocalActivityFactIndex(facts) + let factIndex = factIndex ?? LocalActivityFactIndex(facts) var result: [ActiveTimeWeekEvidence] = [] var workloadMismatchCount = 0 for candidate in usage.sorted( @@ -80,7 +85,10 @@ enum ActiveTimeWeekEvidenceBuilder { continue } let timeline = ActivityTimelineAggregator.evaluate( - facts: factIndex.activityFacts(in: candidate.interval), + facts: factIndex.activityFacts( + in: candidate.interval, + from: facts + ), projections: projections, interval: candidate.interval, observation: scopedObservation( @@ -168,6 +176,10 @@ enum ActiveTimeAvailabilityEngine { static func evaluate( currentUsage: WeeklyUsageEvidence?, activeTimeThisWeek: TimeInterval, + maximumConcurrency: Int = 0, + waitingTime: TimeInterval? = nil, + pollingTime: TimeInterval? = nil, + activityBreakdownReason: String? = nil, activeTimeCoverage: CoverageLevel, activeTimeReason: String?, history: [ActiveTimeWeekEvidence], @@ -177,6 +189,10 @@ enum ActiveTimeAvailabilityEngine { let unavailable: (String) -> ActiveTimeAvailabilitySnapshot = { ActiveTimeAvailabilitySnapshot( activeTimeThisWeek: activeTimeThisWeek, + maximumConcurrency: maximumConcurrency, + waitingTime: waitingTime, + pollingTime: pollingTime, + activityBreakdownReason: activityBreakdownReason, activeTimeCoverage: activeTimeCoverage, activeTimeReason: activeTimeReason, observedInterval: currentUsage?.interval, @@ -272,6 +288,10 @@ enum ActiveTimeAvailabilityEngine { let confidence: ConfidenceLevel = caveats.isEmpty ? .high : .medium return ActiveTimeAvailabilitySnapshot( activeTimeThisWeek: activeTimeThisWeek, + maximumConcurrency: maximumConcurrency, + waitingTime: waitingTime, + pollingTime: pollingTime, + activityBreakdownReason: activityBreakdownReason, activeTimeCoverage: activeTimeCoverage, activeTimeReason: activeTimeReason, observedInterval: currentUsage.interval, diff --git a/Sources/CodexLimits/ActivityTimeline.swift b/Sources/CodexLimits/ActivityTimeline.swift index 4594ec8..f2117e7 100644 --- a/Sources/CodexLimits/ActivityTimeline.swift +++ b/Sources/CodexLimits/ActivityTimeline.swift @@ -59,6 +59,18 @@ struct ActivityTimelineSnapshot: Equatable, Sendable { fileprivate let observation: LocalActivityObservation let interval: DateInterval + func updating( + interval: DateInterval, + observation: LocalActivityObservation + ) -> ActivityTimelineSnapshot { + ActivityTimelineSnapshot( + turns: turns, + boundaryIssues: boundaryIssues, + observation: observation, + interval: interval + ) + } + func slice( in selectedInterval: DateInterval, filters: WorkspaceFilters @@ -304,27 +316,31 @@ enum ActivityTimelineAggregator { start: start, end: end ) - pendingIssues.append( - ( - turnKey, - ActivityTimelineSnapshot.BoundaryIssue( - date: issueDate, - affectedStart: affectedBounds.start, - affectedEnd: affectedBounds.end, - reason: root == nil && taskID != nil - ? "Task Tree metadata is missing" - : "Some Active Turn boundaries are missing or invalid", - rootTaskID: root, - projectLabel: root.flatMap { - projectionsByTask[$0]?.projectLabel - }, - context: context - ) - ) + let issue = ActivityTimelineSnapshot.BoundaryIssue( + date: issueDate, + affectedStart: affectedBounds.start, + affectedEnd: affectedBounds.end, + reason: root == nil && taskID != nil + ? "Task Tree metadata is missing" + : "Some Active Turn boundaries are missing or invalid", + rootTaskID: root, + projectLabel: root.flatMap { + projectionsByTask[$0]?.projectLabel + }, + context: context ) + if issue.affects( + DateInterval( + start: interval.start, + end: .distantFuture + ) + ) { + pendingIssues.append((turnKey, issue)) + } } continue } + guard end > interval.start else { continue } turns.append( ActivityTimelineSnapshot.TurnInterval( start: start, diff --git a/Sources/CodexLimits/AnalyticsWorkspace.swift b/Sources/CodexLimits/AnalyticsWorkspace.swift index d977b6d..5afc9f5 100644 --- a/Sources/CodexLimits/AnalyticsWorkspace.swift +++ b/Sources/CodexLimits/AnalyticsWorkspace.swift @@ -353,10 +353,48 @@ struct UsageChartSelection: Equatable, Sendable { in chart: UsageChartSnapshot, within visibleRange: DateInterval? = nil ) -> UsageChartSelection? { - candidates(in: chart) - .filter { - visibleRange?.contains($0.point.date) ?? true - }.min { + [ + nearestCandidate( + in: chart.allObserved, + series: .observed, + priority: 0, + source: .account, + to: date, + within: visibleRange + ), + nearestCandidate( + in: chart.currentProjection, + series: .currentEstimate, + priority: 1, + source: .derivedEstimate, + to: date, + within: visibleRange + ), + nearestCandidate( + in: chart.historicalProjection, + series: .pastEstimate, + priority: 2, + source: .accountHistory, + to: date, + within: visibleRange + ), + nearestCandidate( + in: chart.estimatedBackfill, + series: .estimatedBackfill, + priority: 3, + source: .tokenEstimate, + to: date, + within: visibleRange + ), + nearestCandidate( + in: chart.target, + series: .target, + priority: 4, + source: .weeklyTarget, + to: date, + within: visibleRange + ) + ].compactMap { $0 }.min { let leftDistance = abs($0.point.date.timeIntervalSince(date)) let rightDistance = abs($1.point.date.timeIntervalSince(date)) if leftDistance == rightDistance { @@ -373,6 +411,30 @@ struct UsageChartSelection: Equatable, Sendable { } } + private static func nearestCandidate( + in points: [UsageChartPoint], + series: UsageChartSeries, + priority: Int, + source: UsageChartPointSource, + to date: Date, + within visibleRange: DateInterval? + ) -> Candidate? { + guard let point = nearestPoint( + in: points, + to: date, + date: \.date, + within: visibleRange + ) else { + return nil + } + return Candidate( + series: series, + point: point, + priority: priority, + source: source + ) + } + private static func candidates( in chart: UsageChartSnapshot ) -> [Candidate] { diff --git a/Sources/CodexLimits/CodexAssistedHistory.swift b/Sources/CodexLimits/CodexAssistedHistory.swift index 0d4b541..f3a6e9b 100644 --- a/Sources/CodexLimits/CodexAssistedHistory.swift +++ b/Sources/CodexLimits/CodexAssistedHistory.swift @@ -44,10 +44,23 @@ actor CodexAssistedHistory { let cutoff: Date } + private enum DeletionMarkerRestore { + case missing + case valid(Date) + case invalid + } + private static let version = 1 + private static let maximumMarkerBytes = 1_048_576 + + private enum HistoryError: Error { + case unreadableHistory + } private let fileURL: URL private let deletionMarkerURL: URL private var file: File? + private var didLoad = false + private var loadFailed = false init( fileURL: URL, @@ -61,7 +74,7 @@ actor CodexAssistedHistory { func results( accountPartitionID: String ) -> [CodexAssistedHistoryResult] { - loadIfNeeded() + guard loadIfNeeded() else { return [] } return file?.results.filter { $0.accountPartitionID == accountPartitionID } ?? [] @@ -70,7 +83,7 @@ actor CodexAssistedHistory { func overhead( accountPartitionID: String ) -> [CodexAssistedHistoryOverhead] { - loadIfNeeded() + guard loadIfNeeded() else { return [] } return file?.overhead.filter { $0.accountPartitionID == accountPartitionID } ?? [] @@ -86,7 +99,9 @@ actor CodexAssistedHistory { guard let accountPartitionID = scope.accountPartitionID else { return } - loadIfNeeded() + guard loadIfNeeded() else { + throw HistoryError.unreadableHistory + } let previous = file file?.overhead.append( CodexAssistedHistoryOverhead( @@ -119,7 +134,20 @@ actor CodexAssistedHistory { } func deleteAll(upTo cutoff: Date = Date()) throws { - loadIfNeeded() + guard loadIfNeeded() else { + try writeDeletionMarker(cutoff) + if FileManager.default.fileExists(atPath: fileURL.path) { + try FileManager.default.removeItem(at: fileURL) + } + if FileManager.default.fileExists( + atPath: deletionMarkerURL.path + ) { + try FileManager.default.removeItem(at: deletionMarkerURL) + } + file = File(version: Self.version, results: [], overhead: []) + loadFailed = false + return + } try writeDeletionMarker(cutoff) file?.results.removeAll { $0.result.observedAt <= cutoff @@ -146,14 +174,34 @@ actor CodexAssistedHistory { } } - private func loadIfNeeded() { - guard file == nil else { return } - let deletionCutoff = deletionMarker() - guard let data = try? Data(contentsOf: fileURL), + @discardableResult + private func loadIfNeeded() -> Bool { + if didLoad { return !loadFailed } + didLoad = true + let deletionCutoff: Date? + switch deletionMarker() { + case .missing: + deletionCutoff = nil + case let .valid(cutoff): + deletionCutoff = cutoff + case .invalid: + loadFailed = true + return false + } + guard FileManager.default.fileExists(atPath: fileURL.path) else { + file = File(version: Self.version, results: [], overhead: []) + return true + } + // ponytail: v1 is one JSON document; map its bytes until a + // future file version can decode records one at a time. + guard let data = try? Data( + contentsOf: fileURL, + options: .mappedIfSafe + ), let decoded = try? JSONDecoder().decode(File.self, from: data), decoded.version == Self.version else { - file = File(version: Self.version, results: [], overhead: []) - return + loadFailed = true + return false } var filtered = decoded if let deletionCutoff { @@ -165,6 +213,7 @@ actor CodexAssistedHistory { } } file = filtered + return true } private func persist() throws { @@ -189,15 +238,40 @@ actor CodexAssistedHistory { try data.write(to: deletionMarkerURL, options: .atomic) } - private func deletionMarker() -> Date? { - guard let data = try? Data(contentsOf: deletionMarkerURL), + private func deletionMarker() -> DeletionMarkerRestore { + guard FileManager.default.fileExists( + atPath: deletionMarkerURL.path + ) else { + return .missing + } + guard let data = Self.readData( + at: deletionMarkerURL, + maximumBytes: Self.maximumMarkerBytes + ), let marker = try? JSONDecoder().decode( DeletionMarker.self, from: data ) else { + return .invalid + } + return .valid(marker.cutoff) + } + + private static func readData( + at url: URL, + maximumBytes: Int + ) -> Data? { + guard let handle = try? FileHandle(forReadingFrom: url) else { + return nil + } + defer { try? handle.close() } + guard let data = try? handle.read( + upToCount: maximumBytes + 1 + ), + data.count <= maximumBytes else { return nil } - return marker.cutoff + return data } private static func defaultFileURL() -> URL { diff --git a/Sources/CodexLimits/CodexAssistedInsights.swift b/Sources/CodexLimits/CodexAssistedInsights.swift index 5a6411a..8d09a2a 100644 --- a/Sources/CodexLimits/CodexAssistedInsights.swift +++ b/Sources/CodexLimits/CodexAssistedInsights.swift @@ -204,16 +204,10 @@ struct CodexMetadataAnalysisPayload: Codable, Equatable, Sendable { if exploration.filters.isEmpty { local = reader.localTokenActivity.slice(in: localRange) } else { - let receiptSlice = reader.usageReceipts.slice( + local = reader.usageReceipts.localTokenSlice( in: localRange, filters: exploration.filters ) - local = LocalTokenActivitySlice( - tokens: receiptSlice.totalTokens, - points: receiptSlice.points, - coverage: receiptSlice.coverage, - reason: receiptSlice.reason - ) } let activity = reader.activityTimeline.slice( in: localRange, diff --git a/Sources/CodexLimits/CodexClient.swift b/Sources/CodexLimits/CodexClient.swift index be3c0df..aaad5d6 100644 --- a/Sources/CodexLimits/CodexClient.swift +++ b/Sources/CodexLimits/CodexClient.swift @@ -69,15 +69,19 @@ enum CodexClientError: LocalizedError { } final class CodexAppServerConnection: @unchecked Sendable { + private static let defaultMaximumLineBytes = 16 * 1_024 * 1_024 + let input: FileHandle let isRunning: () -> Bool let stop: () -> Void private let outputDescriptor: Int32 + private let maximumLineBytes: Int private var bufferedOutput = Data() init( input: FileHandle, output: FileHandle, + maximumLineBytes: Int = defaultMaximumLineBytes, isRunning: @escaping () -> Bool, stop: @escaping () -> Void ) { @@ -85,6 +89,7 @@ final class CodexAppServerConnection: @unchecked Sendable { self.isRunning = isRunning self.stop = stop outputDescriptor = Darwin.dup(output.fileDescriptor) + self.maximumLineBytes = max(maximumLineBytes, 1) } deinit { @@ -94,12 +99,26 @@ final class CodexAppServerConnection: @unchecked Sendable { } func readLine() async -> Data? { + var searchedByteCount = 0 while true { - if let newline = bufferedOutput.firstIndex(of: 0x0A) { + let searchStart = bufferedOutput.index( + bufferedOutput.startIndex, + offsetBy: min(searchedByteCount, bufferedOutput.count) + ) + if let newline = bufferedOutput[searchStart...].firstIndex( + of: 0x0A + ) { + guard newline <= maximumLineBytes else { + return closeOversizedLine() + } let line = bufferedOutput[.. maximumLineBytes { + return closeOversizedLine() + } let chunk = await Task.detached { [outputDescriptor] () -> Data? in var data = Data(count: 64 * 1_024) @@ -122,6 +141,12 @@ final class CodexAppServerConnection: @unchecked Sendable { bufferedOutput.append(chunk) } } + + private func closeOversizedLine() -> Data? { + bufferedOutput = Data() + stop() + return nil + } } enum CodexIsolatedHome { @@ -329,6 +354,7 @@ actor CodexClient { ) async throws -> T { await protocolGate.enter() do { + try Task.checkCancellation() let result = try await operation() await protocolGate.leave() return result @@ -740,7 +766,7 @@ actor CodexClient { let snapshots = result.rateLimitsByLimitId ?? ["codex": result.rateLimits] let mainSnapshot = snapshots["codex"] ?? result.rateLimits - return windows(from: mainSnapshot) + return try windows(from: mainSnapshot) .first(where: { $0.durationMinutes == weeklyWindowDurationMinutes }) @@ -777,7 +803,7 @@ actor CodexClient { let snapshots = rateResult.rateLimitsByLimitId ?? ["codex": rateResult.rateLimits] let mainSnapshot = snapshots["codex"] ?? rateResult.rateLimits - let mainWindows = windows(from: mainSnapshot) + let mainWindows = try windows(from: mainSnapshot) let mainWindow = mainWindows.first(where: { $0.durationMinutes == weeklyWindowDurationMinutes }) @@ -787,10 +813,10 @@ actor CodexClient { .map { LimitReading(limitId: "codex", name: windowName($0.durationMinutes), window: $0) } - let otherLimits = snapshots + let otherLimits = try snapshots .filter { $0.key != "codex" } .compactMap { id, snapshot -> LimitReading? in - guard let window = windows(from: snapshot).min(by: { + guard let window = try windows(from: snapshot).min(by: { $0.remainingPercent < $1.remainingPercent }) else { return nil } return LimitReading( @@ -847,6 +873,9 @@ actor CodexClient { reached: mainSnapshot.spendControlReached ) } + guard spendControl?.isValid != false else { + throw CodexClientError.invalidResponse + } let facts = AccountFacts( lifetimeTokens: summary?.lifetimeTokens, peakDailyTokens: summary?.peakDailyTokens, @@ -900,7 +929,7 @@ actor CodexClient { ) } - return UsageSnapshot( + let snapshot = UsageSnapshot( mainLimit: mainWindow.map { LimitReading(limitId: "codex", name: "Codex", window: $0) }, @@ -912,6 +941,10 @@ actor CodexClient { fetchedAt: fetchedAt, accountFacts: facts.isEmpty ? nil : facts ) + guard snapshot.isValid else { + throw CodexClientError.invalidResponse + } + return snapshot } static func decodeAccount(_ response: Data) throws -> CodexAccountObservation { @@ -988,11 +1021,18 @@ actor CodexClient { ) } - private static func windows(from snapshot: RateLimitSnapshot) -> [UsageWindow] { - [snapshot.primary, snapshot.secondary].compactMap { window in - guard let window, - let resetsAt = window.resetsAt, - let duration = window.windowDurationMins else { return nil } + private static func windows( + from snapshot: RateLimitSnapshot + ) throws -> [UsageWindow] { + try [snapshot.primary, snapshot.secondary].compactMap { window in + guard let window else { return nil } + guard let resetsAt = window.resetsAt, + let duration = window.windowDurationMins else { + return nil + } + guard window.usedPercent.isFinite, duration > 0 else { + throw CodexClientError.invalidResponse + } return UsageWindow( remainingPercent: min(max(100 - window.usedPercent, 0), 100), resetsAt: Date(timeIntervalSince1970: TimeInterval(resetsAt)), diff --git a/Sources/CodexLimits/ForecastEngine.swift b/Sources/CodexLimits/ForecastEngine.swift index 025f99a..7bbc64c 100644 --- a/Sources/CodexLimits/ForecastEngine.swift +++ b/Sources/CodexLimits/ForecastEngine.swift @@ -152,26 +152,49 @@ enum ForecastEngine { now: Date ) -> Double? { let day: TimeInterval = 86_400 - let dayNumber: (Date) -> Int = { Int(floor($0.timeIntervalSince1970 / day)) } - let start = dayNumber(window.startsAt) - let today = dayNumber(now) - let buckets = Dictionary(grouping: tokenHistory, by: { dayNumber($0.date) }) - .mapValues { $0.reduce(Int64(0)) { $0 + $1.tokens } } + let dayNumber: (Date) -> Int? = { + let value = floor($0.timeIntervalSince1970 / day) + guard value.isFinite else { return nil } + return Int(exactly: value) + } + guard let start = dayNumber(window.startsAt), + let today = dayNumber(now) else { return nil } + var buckets: [Int: Double] = [:] + for tokenDay in tokenHistory { + guard tokenDay.tokens >= 0, + let bucket = dayNumber(tokenDay.date) else { return nil } + let total = (buckets[bucket] ?? 0) + Double(tokenDay.tokens) + guard total.isFinite else { return nil } + buckets[bucket] = total + } guard let first = buckets.keys.min(), let latest = buckets.keys.filter({ $0 < today }).max(), latest >= start else { return nil } - let currentCount = latest - start + 1 - let currentTokens = (start ... latest).reduce(Int64(0)) { $0 + (buckets[$1] ?? 0) } - let historyEnd = start - 1 - let historyStart = max(first, historyEnd - 27) + let currentCount = Double(latest) - Double(start) + 1 + let currentTokens = buckets.reduce(0.0) { + $1.key >= start && $1.key <= latest ? $0 + $1.value : $0 + } + let (historyEnd, endOverflow) = start.subtractingReportingOverflow(1) + guard !endOverflow else { return nil } + let (candidateStart, startOverflow) = + historyEnd.subtractingReportingOverflow(27) + guard !startOverflow else { return nil } + let historyStart = max(first, candidateStart) guard historyStart <= historyEnd, currentTokens > 0 else { return nil } - let historyCount = historyEnd - historyStart + 1 - let historyTokens = (historyStart ... historyEnd).reduce(Int64(0)) { $0 + (buckets[$1] ?? 0) } - let currentAverage = Double(currentTokens) / Double(currentCount) - let historicalAverage = Double(historyTokens) / Double(historyCount) - guard currentAverage > 0, historicalAverage > 0 else { return nil } + let historyCount = Double(historyEnd) - Double(historyStart) + 1 + let historyTokens = buckets.reduce(0.0) { + $1.key >= historyStart && $1.key <= historyEnd + ? $0 + $1.value + : $0 + } + let currentAverage = currentTokens / currentCount + let historicalAverage = historyTokens / historyCount + guard currentAverage.isFinite, + historicalAverage.isFinite, + currentAverage > 0, + historicalAverage > 0 else { return nil } // Daily token buckets are a coarse bootstrap; percentage-based windows replace them. let relativePace = min(max(historicalAverage / currentAverage, 0.25), 4) diff --git a/Sources/CodexLimits/LocalActivityCollector.swift b/Sources/CodexLimits/LocalActivityCollector.swift index 65824c9..4b81c80 100644 --- a/Sources/CodexLimits/LocalActivityCollector.swift +++ b/Sources/CodexLimits/LocalActivityCollector.swift @@ -1,22 +1,29 @@ import CryptoKit import Foundation +private func nextRevision(after revision: UInt64) -> UInt64 { + revision == .max ? 1 : revision + 1 +} + struct LocalActivityCollection: Equatable, Sendable { let facts: [LocalActivityFact] let projections: [ThreadProjection] let observation: LocalActivityObservation let bytesRead: UInt64 + let contentRevision: UInt64 static func unavailable( _ reason: String, facts: [LocalActivityFact] = [], - projections: [ThreadProjection] = [] + projections: [ThreadProjection] = [], + contentRevision: UInt64 = 0 ) -> LocalActivityCollection { LocalActivityCollection( facts: facts, projections: projections, observation: .unavailable(reason), - bytesRead: 0 + bytesRead: 0, + contentRevision: contentRevision ) } @@ -32,7 +39,8 @@ struct LocalActivityCollection: Equatable, Sendable { observedAt: observedAt, reason: reason ), - bytesRead: bytesRead + bytesRead: bytesRead, + contentRevision: contentRevision ) case .unavailable: return self @@ -41,6 +49,17 @@ struct LocalActivityCollection: Equatable, Sendable { } actor LocalActivityCollector { + private static let maximumMetadataBytes = 1_048_576 + private static let maximumRefreshLines = 10_000 + private static let maximumRefreshBytes: UInt64 = 8 * 1_024 * 1_024 + private static let maximumRolloutRecordBytes = 7 * 1_024 * 1_024 + + private struct ObservationSignature: Equatable { + let sourceVersion: String? + let reason: String? + let coverage: CoverageLevel + } + private struct FileState { var cursor: RolloutCursor var normalization: LocalActivityNormalizationState @@ -53,6 +72,20 @@ actor LocalActivityCollector { var storageFingerprint: String? var factsLoaded: Bool var requiresContextRebuild: Bool + var factRestoreOffset: UInt64 + var factRestoreFileSize: UInt64? + var restoredFactIdentities: Set + } + + private struct FactIdentity: Hashable { + let eventID: String + let key: String + } + + private enum FactLoadOutcome { + case ready(linesRead: Int, bytesRead: UInt64) + case partial(linesRead: Int, bytesRead: UInt64) + case invalid } private struct PersistedFile: Codable { @@ -103,11 +136,25 @@ actor LocalActivityCollector { private var nextProjectionListCursor: String? private var hasStartedProjectionList = false private var hasCompleteProjectionList = false + private var lastProjectionListSucceeded: Bool? + private var attemptedProjectionTaskIDs = Set() private var stateGeneration: UInt64 = 0 private var historyCutoff: Date? private var historyDeletionPending = false private var deletionMarkerInvalid = false private var restoreWarning: String? + private var cachedFacts: [LocalActivityFact]? + private var cachedFactPaths = Set() + private var cachedFactsIntervalStart: Date? + private var lastPublishedProjectionIdentities: [ProjectionIdentity]? + private var lastObservationSignature: ObservationSignature? + private var contentRevision: UInt64 = 0 + private var importContinuationPending = false + private var pendingFactRestorePaths = Set() + private var publishedFactPaths = Set() + private var refreshGeneration: UInt64 = 0 + private var didReadInstalledCLIVersion = false + private var cachedInstalledCLIVersion: String? init( rootDirectory: URL = FileManager.default.homeDirectoryForCurrentUser @@ -145,7 +192,7 @@ actor LocalActivityCollector { func selectPartition(_ id: String) { guard partitionID != id else { return } - stateGeneration &+= 1 + stateGeneration = nextRevision(after: stateGeneration) persist() partitionID = id files.removeAll() @@ -158,14 +205,18 @@ actor LocalActivityCollector { hasStartedProjectionList = false hasCompleteProjectionList = false restoreWarning = nil + clearPublishedContent() restore() } func refresh( interval: DateInterval, - observedAt: Date = Date() + observedAt: Date = Date(), + refreshMetadata: Bool = true ) async -> LocalActivityCollection { - let refreshGeneration = stateGeneration + refreshGeneration = nextRevision(after: refreshGeneration) + let currentRefreshGeneration = refreshGeneration + let currentStateGeneration = stateGeneration if historyDeletionPending { return .unavailable( deletionMarkerInvalid @@ -176,48 +227,165 @@ actor LocalActivityCollector { guard FileManager.default.fileExists(atPath: rootDirectory.path) else { return .unavailable( "Codex local records are unavailable", - facts: files.values.flatMap(\.facts) + facts: cachedFacts ?? files.values.flatMap(\.facts) ) } let projectionsBeforeRefresh = projections - let listSucceeded = await refreshProjectionList( - generation: refreshGeneration - ) - guard refreshGeneration == stateGeneration else { + let listSucceeded: Bool? + if refreshMetadata { + listSucceeded = await refreshProjectionList( + generation: currentStateGeneration, + refreshGeneration: currentRefreshGeneration + ) + } else { + listSucceeded = lastProjectionListSucceeded + } + guard !Task.isCancelled else { + return .unavailable("Local activity read was cancelled") + } + guard currentStateGeneration == stateGeneration, + currentRefreshGeneration == refreshGeneration else { return .unavailable("Account changed during local activity read") } - let calendarFiles = rolloutFiles(in: interval) + if refreshMetadata { + lastProjectionListSucceeded = listSucceeded + } let currentProjections = activeProjections(in: interval) - let projectedFiles = Set( - currentProjections.compactMap(validProjectionFile) - ) - let trackedFiles = Set(files.compactMap { path, state in - stateHasActivity(state, in: interval) ? fileURL(path) : nil - }) - let candidatePaths = Set( - calendarFiles - .union(projectedFiles) - .union(trackedFiles) - .compactMap(fileKey) - ) + let candidatePaths: Set + if !refreshMetadata, + cachedFactsIntervalStart == interval.start, + !cachedFactPaths.isEmpty { + candidatePaths = cachedFactPaths + } else { + let calendarFiles = rolloutFiles(in: interval) + let projectedFiles = Set( + currentProjections.compactMap(validProjectionFile) + ) + let trackedFiles = Set(files.compactMap { path, state in + stateHasActivity(state, in: interval) ? fileURL(path) : nil + }) + candidatePaths = Set( + calendarFiles + .union(projectedFiles) + .union(trackedFiles) + .compactMap(fileKey) + ) + } + for path in candidatePaths { + bindRestoredStateIfNeeded(to: path) + } + if cachedFactsIntervalStart != interval.start { + restartPartialFactRestores() + } + let reusesPublishedFacts = cachedFacts != nil + && cachedFactsIntervalStart == interval.start + && cachedFactPaths.isSubset(of: candidatePaths) + if !reusesPublishedFacts { + pendingFactRestorePaths = Set(candidatePaths.filter { + files[$0]?.factsLoaded == false + }) + publishedFactPaths.removeAll() + } else { + pendingFactRestorePaths.formIntersection(candidatePaths) + for path in candidatePaths.subtracting(cachedFactPaths) + where files[path]?.factsLoaded == false { + pendingFactRestorePaths.insert(path) + } + } var bytesRead: UInt64 = 0 + var factsChanged = false + var rewroteFacts = false + var appendedFacts: [LocalActivityFact] = [] + var remainingLineBudget = Self.maximumRefreshLines + var remainingByteBudget = Self.maximumRefreshBytes + var importStillInProgress = false var gapReason = restoreWarning ?? (listSucceeded == false ? "Local task discovery is incomplete" : currentProjections.contains { validProjectionFile($0) == nil } ? "Local rollout path is unavailable" : nil) - for path in candidatePaths.sorted() { - bindRestoredStateIfNeeded(to: path) + for path in candidatePaths.sorted(by: >) { + guard !Task.isCancelled else { + return .unavailable("Local activity read was cancelled") + } + if remainingLineBudget == 0 || remainingByteBudget == 0 { + importStillInProgress = true + gapReason = gapReason + ?? "Local task import is still in progress" + break + } let file = fileURL(path) - loadFactsIfNeeded(for: path) + let pathWasPublished = publishedFactPaths.contains(path) + if !reusesPublishedFacts + || pendingFactRestorePaths.contains(path) { + let factsBeforeRestore = files[path]?.facts.count ?? 0 + let factLoad = loadFactsIfNeeded( + for: path, + interval: interval, + maximumLines: remainingLineBudget, + maximumBytes: remainingByteBudget + ) + switch factLoad { + case let .ready(linesRead, factBytesRead): + pendingFactRestorePaths.remove(path) + remainingLineBudget = max( + remainingLineBudget - linesRead, + 0 + ) + remainingByteBudget = factBytesRead + >= remainingByteBudget + ? 0 + : remainingByteBudget - factBytesRead + case let .partial(linesRead, factBytesRead): + importStillInProgress = true + remainingLineBudget = max( + remainingLineBudget - linesRead, + 0 + ) + remainingByteBudget = factBytesRead + >= remainingByteBudget + ? 0 + : remainingByteBudget - factBytesRead + gapReason = gapReason + ?? "Local task import is still in progress" + if reusesPublishedFacts, + let facts = files[path]?.facts, + facts.count > factsBeforeRestore { + appendedFacts.append( + contentsOf: facts[factsBeforeRestore...] + ) + } + publishedFactPaths.insert(path) + continue + case .invalid: + pendingFactRestorePaths.remove(path) + break + } + if reusesPublishedFacts, + let facts = files[path]?.facts, + facts.count > factsBeforeRestore { + appendedFacts.append( + contentsOf: facts[factsBeforeRestore...] + ) + } + if files[path] != nil { + publishedFactPaths.insert(path) + } + if remainingLineBudget == 0 || remainingByteBudget == 0 { + importStillInProgress = true + gapReason = gapReason + ?? "Local task import is still in progress" + continue + } + } guard FileManager.default.fileExists(atPath: file.path) else { gapReason = gapReason ?? "Local task records are missing" continue } + var previous = files.removeValue(forKey: path) do { - let previous = files[path] if previous?.hasMalformedRecords == true { gapReason = gapReason ?? "Some local diagnostic records could not be read" @@ -229,14 +397,38 @@ actor LocalActivityCollector { cursor: requiresContextRebuild ? nil : previous?.cursor, - observedAt: observedAt + observedAt: observedAt, + maximumLines: remainingLineBudget, + maximumBytes: remainingByteBudget, + maximumRecordBytes: Self.maximumRolloutRecordBytes + ) + remainingLineBudget = max( + remainingLineBudget - batch.processedLineCount, + 0 ) - bytesRead += batch.bytesRead + remainingByteBudget = batch.bytesRead + >= remainingByteBudget + ? 0 + : remainingByteBudget - batch.bytesRead + let (totalBytesRead, bytesOverflowed) = + bytesRead.addingReportingOverflow(batch.bytesRead) + if bytesOverflowed { + bytesRead = .max + gapReason = gapReason + ?? "Local task record size is invalid" + } else { + bytesRead = totalBytesRead + } if batch.malformedRecordCount > 0 { gapReason = gapReason ?? "Some local diagnostic records could not be read" } - if batch.requiresRebuild, previous != nil { + if batch.hasMoreRecords { + importStillInProgress = true + gapReason = gapReason + ?? "Local task import is still in progress" + } + if batch.continuityChanged, previous != nil { gapReason = gapReason ?? "Local task record continuity changed" } @@ -245,14 +437,15 @@ actor LocalActivityCollector { previous != nil, !requiresContextRebuild { if batch.malformedRecordCount > 0 { - files[path]?.hasMalformedRecords = true + previous?.hasMalformedRecords = true } if previous?.cursor != batch.cursor { - files[path]?.cursor = batch.cursor + previous?.cursor = batch.cursor changedPaths.insert(path) } else if batch.malformedRecordCount > 0 { changedPaths.insert(path) } + files[path] = previous continue } let normalized = normalizer.normalize( @@ -267,6 +460,8 @@ actor LocalActivityCollector { let rewritesFacts = batch.requiresRebuild || requiresContextRebuild || previous == nil + let rewritesPublishedFacts = rewritesFacts + && pathWasPublished let newFacts: [LocalActivityFact] if rewritesFacts { newFacts = factsAfterHistoryCutoff(normalized.facts) @@ -279,31 +474,67 @@ actor LocalActivityCollector { return !existingEventIDs.contains(eventID) } } - let combinedFacts = rewritesFacts - ? newFacts - : (previous?.facts ?? []) + newFacts - let activityBounds = tokenActivityBounds(combinedFacts) - let discontinuityAt = batch.requiresRebuild + let discontinuityAt = batch.continuityChanged && previous != nil && !requiresContextRebuild ? observedAt : previous?.discontinuityAt - files[path] = FileState( - cursor: batch.cursor, - normalization: normalized.state, - facts: combinedFacts, - eventIDs: Set(combinedFacts.compactMap(\.eventID)), - activityStart: activityBounds?.start, - activityEnd: activityBounds?.end, - discontinuityAt: discontinuityAt, - hasMalformedRecords: batch.requiresRebuild - ? batch.malformedRecordCount > 0 - : previous?.hasMalformedRecords == true - || batch.malformedRecordCount > 0, - storageFingerprint: previous?.storageFingerprint, - factsLoaded: true, - requiresContextRebuild: false - ) + let hasMalformedRecords = batch.requiresRebuild + ? batch.malformedRecordCount > 0 + : previous?.hasMalformedRecords == true + || batch.malformedRecordCount > 0 + if rewritesFacts { + let activityBounds = tokenActivityBounds(newFacts) + previous = FileState( + cursor: batch.cursor, + normalization: normalized.state, + facts: newFacts, + eventIDs: Set(newFacts.compactMap(\.eventID)), + activityStart: activityBounds?.start, + activityEnd: activityBounds?.end, + discontinuityAt: discontinuityAt, + hasMalformedRecords: hasMalformedRecords, + storageFingerprint: previous?.storageFingerprint, + factsLoaded: true, + requiresContextRebuild: false, + factRestoreOffset: 0, + factRestoreFileSize: nil, + restoredFactIdentities: [] + ) + } else { + let hadAllFacts = previous?.factsLoaded == true + previous?.cursor = batch.cursor + previous?.normalization = normalized.state + if hadAllFacts { + previous?.facts.append(contentsOf: newFacts) + } + previous?.eventIDs.formUnion( + newFacts.compactMap(\.eventID) + ) + if let bounds = tokenActivityBounds(newFacts) { + let activityStart = previous?.activityStart + let activityEnd = previous?.activityEnd + previous?.activityStart = activityStart + .map { min($0, bounds.start) } + ?? bounds.start + previous?.activityEnd = activityEnd + .map { max($0, bounds.end) } + ?? bounds.end + } + previous?.discontinuityAt = discontinuityAt + previous?.hasMalformedRecords = hasMalformedRecords + previous?.factsLoaded = hadAllFacts + previous?.requiresContextRebuild = false + } + files[path] = previous + publishedFactPaths.insert(path) + factsChanged = factsChanged + || rewritesFacts + || !newFacts.isEmpty + rewroteFacts = rewroteFacts || rewritesPublishedFacts + if !rewritesPublishedFacts { + appendedFacts.append(contentsOf: newFacts) + } changedPaths.insert(path) scheduleFactWrite( path: path, @@ -311,6 +542,7 @@ actor LocalActivityCollector { rewritesFile: rewritesFacts ) } catch { + files[path] = previous gapReason = gapReason ?? "Local task records are missing" } } @@ -323,7 +555,7 @@ actor LocalActivityCollector { return discontinuityAt >= interval.start && discontinuityAt <= interval.end }) { - gapReason = gapReason ?? "Local task record continuity changed" + gapReason = "Local task record continuity changed" } let activeTaskIDs = taskIDs(in: activeStates) let projectionChainsBefore = projectionChainIdentities( @@ -332,24 +564,27 @@ actor LocalActivityCollector { ) if await completeProjections( for: activeTaskIDs, - listSucceeded: listSucceeded, - generation: refreshGeneration + listSucceeded: refreshMetadata ? listSucceeded : true, + generation: currentStateGeneration, + refreshGeneration: currentRefreshGeneration, + retriesMissingProjections: refreshMetadata ) == false { gapReason = gapReason ?? "Local task metadata is incomplete" } + guard !Task.isCancelled else { + return .unavailable("Local activity read was cancelled") + } if projectionChainsBefore != projectionChainIdentities( for: activeTaskIDs ) { markFilesChanged(for: activeTaskIDs) } - guard refreshGeneration == stateGeneration else { + guard currentStateGeneration == stateGeneration, + currentRefreshGeneration == refreshGeneration else { return .unavailable("Account changed during local activity read") } - if activeStates.contains(where: { state in - state.facts.contains { $0.key == .token } - && !state.facts.contains { - $0.key == .task && $0.availability == .available - } + if activeStates.contains(where: { + $0.activityStart != nil && taskID(in: $0) == nil }) { gapReason = gapReason ?? "Local task identity is missing" } @@ -358,12 +593,27 @@ actor LocalActivityCollector { .map(\.normalization.sourceVersion) .filter { $0 != "unknown" } ) - let hasTokenFacts = activeStates.contains { - $0.facts.contains { $0.key == .token } - } + let hasTokenFacts = activeStates.contains { $0.activityStart != nil } if hasTokenFacts { - let installedVersion = await installedCLIVersion?() - guard refreshGeneration == stateGeneration else { + let installedVersion: String? + if refreshMetadata || !didReadInstalledCLIVersion { + installedVersion = await installedCLIVersion?() + guard !Task.isCancelled else { + return .unavailable("Local activity read was cancelled") + } + guard currentStateGeneration == stateGeneration, + currentRefreshGeneration == refreshGeneration else { + return .unavailable( + "Account changed during local activity read" + ) + } + cachedInstalledCLIVersion = installedVersion + didReadInstalledCLIVersion = true + } else { + installedVersion = cachedInstalledCLIVersion + } + guard currentStateGeneration == stateGeneration, + currentRefreshGeneration == refreshGeneration else { return .unavailable("Account changed during local activity read") } if versions.isEmpty { @@ -404,26 +654,168 @@ actor LocalActivityCollector { }.sorted { $0.taskID < $1.taskID }.map(\.withoutRolloutFileURL) - if !persist() { + let activeProjectionIdentities = activeProjections.map { + ProjectionIdentity( + taskID: $0.taskID, + parentTaskID: $0.parentTaskID, + projectLabel: $0.projectLabel, + createdAt: $0.createdAt, + updatedAt: $0.updatedAt + ) + } + guard currentStateGeneration == stateGeneration, + currentRefreshGeneration == refreshGeneration else { + return .unavailable("Account changed during local activity read") + } + guard !Task.isCancelled else { + return .unavailable("Local activity read was cancelled") + } + let persisted = persist() + projections = projections.filter { + activeProjectionIDs.contains($0.key) + } + attemptedProjectionTaskIDs.formIntersection(activeProjectionIDs) + if !persisted { gapReason = gapReason ?? "Local activity could not be saved" } - return LocalActivityCollection( - facts: activeStates.flatMap(\.facts), - projections: activeProjections, - observation: gapReason.map { - .gap( - sourceVersion: version, - observedAt: observedAt, - reason: $0 + let factsWereRebuilt: Bool + let mergedFacts: [LocalActivityFact] + if reusesPublishedFacts, !rewroteFacts { + var currentFacts = cachedFacts ?? [] + cachedFacts = nil + if !appendedFacts.isEmpty { + currentFacts.append( + contentsOf: factsForActiveInterval( + appendedFacts, + interval: interval + ) ) - } ?? .continuous( + } + cachedFacts = currentFacts + mergedFacts = currentFacts + factsWereRebuilt = false + } else { + if rewroteFacts { + pendingFactRestorePaths.formUnion( + candidatePaths.filter { + files[$0]?.factsLoaded == false + } + ) + for path in candidatePaths.sorted(by: >) + where files[path]?.factsLoaded == false { + guard !Task.isCancelled else { + return .unavailable( + "Local activity read was cancelled" + ) + } + guard remainingLineBudget > 0, + remainingByteBudget > 0 else { + importStillInProgress = true + gapReason = gapReason + ?? "Local task import is still in progress" + continue + } + let outcome = loadFactsIfNeeded( + for: path, + interval: interval, + maximumLines: remainingLineBudget, + maximumBytes: remainingByteBudget + ) + let read: (Int, UInt64) + switch outcome { + case let .ready(linesRead, factBytesRead), + let .partial(linesRead, factBytesRead): + read = (linesRead, factBytesRead) + case .invalid: + pendingFactRestorePaths.remove(path) + gapReason = gapReason + ?? "Saved local activity could not be read" + continue + } + remainingLineBudget = max( + remainingLineBudget - read.0, + 0 + ) + remainingByteBudget = read.1 >= remainingByteBudget + ? 0 + : remainingByteBudget - read.1 + if case .partial = outcome { + importStillInProgress = true + gapReason = gapReason + ?? "Local task import is still in progress" + } else { + pendingFactRestorePaths.remove(path) + } + } + } + var rebuiltFacts: [LocalActivityFact] = [] + for path in candidatePaths.sorted() { + guard let facts = files[path]?.facts else { continue } + rebuiltFacts.append( + contentsOf: facts.lazy.filter { + self.isFactActive($0, in: interval) + } + ) + } + mergedFacts = rebuiltFacts + cachedFacts = mergedFacts + publishedFactPaths = Set(candidatePaths.filter { + guard let state = files[$0] else { return false } + return state.factsLoaded || state.factRestoreOffset > 0 + }) + cachedFactPaths = candidatePaths + cachedFactsIntervalStart = interval.start + factsWereRebuilt = true + } + if reusesPublishedFacts, !rewroteFacts { + cachedFactPaths = candidatePaths + cachedFactsIntervalStart = interval.start + } + let observation: LocalActivityObservation = gapReason.map { + .gap( sourceVersion: version, - observedAt: observedAt - ), - bytesRead: bytesRead + observedAt: observedAt, + reason: $0 + ) + } ?? .continuous( + sourceVersion: version, + observedAt: observedAt + ) + let observationSignature = ObservationSignature( + sourceVersion: version == "unknown" ? nil : version, + reason: observation.reason, + coverage: observation.coverage + ) + if factsWereRebuilt + || factsChanged + || !appendedFacts.isEmpty + || lastPublishedProjectionIdentities != activeProjectionIdentities + || lastObservationSignature != observationSignature { + advanceContentRevision() + } + lastPublishedProjectionIdentities = activeProjectionIdentities + lastObservationSignature = observationSignature + if persisted { + unloadPersistedFacts(activePaths: candidatePaths) + } + importContinuationPending = importStillInProgress + return LocalActivityCollection( + facts: mergedFacts, + projections: activeProjections, + observation: observation, + bytesRead: bytesRead, + contentRevision: contentRevision ) } + private func advanceContentRevision() { + contentRevision = nextRevision(after: contentRevision) + } + + func hasPendingImport() -> Bool { + importContinuationPending + } + func deleteHistory(at deletedAt: Date = Date()) throws { historyCutoff = deletedAt historyDeletionPending = stateDirectory != nil @@ -435,7 +827,7 @@ actor LocalActivityCollector { for: stateDirectory ) } - stateGeneration &+= 1 + stateGeneration = nextRevision(after: stateGeneration) files.removeAll() restoredFilesByFingerprint.removeAll() restoredFingerprintByIdentity.removeAll() @@ -446,6 +838,7 @@ actor LocalActivityCollector { hasStartedProjectionList = false hasCompleteProjectionList = false restoreWarning = nil + clearPublishedContent() guard let stateDirectory else { return } if FileManager.default.fileExists(atPath: stateDirectory.path) { try FileManager.default.removeItem(at: stateDirectory) @@ -479,7 +872,7 @@ actor LocalActivityCollector { ) historyCutoff = cutoff deletionMarkerInvalid = false - stateGeneration &+= 1 + stateGeneration = nextRevision(after: stateGeneration) files.removeAll() restoredFilesByFingerprint.removeAll() restoredFingerprintByIdentity.removeAll() @@ -490,6 +883,7 @@ actor LocalActivityCollector { hasStartedProjectionList = false hasCompleteProjectionList = false restoreWarning = nil + clearPublishedContent() if FileManager.default.fileExists(atPath: stateDirectory.path) { try FileManager.default.removeItem(at: stateDirectory) } @@ -515,7 +909,7 @@ actor LocalActivityCollector { try FileManager.default.removeItem(at: markerURL) } } - stateGeneration &+= 1 + stateGeneration = nextRevision(after: stateGeneration) files.removeAll() restoredFilesByFingerprint.removeAll() restoredFingerprintByIdentity.removeAll() @@ -526,6 +920,7 @@ actor LocalActivityCollector { hasStartedProjectionList = false hasCompleteProjectionList = false restoreWarning = nil + clearPublishedContent() historyCutoff = nil historyDeletionPending = false deletionMarkerInvalid = false @@ -577,7 +972,10 @@ actor LocalActivityCollector { Set(states.compactMap(taskID(in:))) } - private func refreshProjectionList(generation: UInt64) async -> Bool? { + private func refreshProjectionList( + generation: UInt64, + refreshGeneration: UInt64 + ) async -> Bool? { guard let projectionSource else { return nil } let hadStarted = hasStartedProjectionList do { @@ -585,7 +983,10 @@ actor LocalActivityCollector { cursor: nil, limit: 100 ) - guard generation == stateGeneration else { return nil } + guard generation == stateGeneration, + refreshGeneration == self.refreshGeneration else { + return nil + } for projection in newestPage.tasks { projections[projection.taskID] = projection } @@ -612,7 +1013,10 @@ actor LocalActivityCollector { cursor: cursor, limit: 100 ) - guard generation == stateGeneration else { return nil } + guard generation == stateGeneration, + refreshGeneration == self.refreshGeneration else { + return nil + } for projection in page.tasks { projections[projection.taskID] = projection } @@ -632,7 +1036,9 @@ actor LocalActivityCollector { private func completeProjections( for taskIDs: Set, listSucceeded: Bool?, - generation: UInt64 + generation: UInt64, + refreshGeneration: UInt64, + retriesMissingProjections: Bool ) async -> Bool { guard let projectionSource else { return true } guard !taskIDs.isEmpty else { return true } @@ -643,19 +1049,30 @@ actor LocalActivityCollector { while let taskID = pending.first { pending.removeFirst() guard inspected.insert(taskID).inserted else { continue } - let mustRead = projections[taskID] == nil - || (taskIDs.contains(taskID) && listSucceeded != true) + let mustRead = ( + projections[taskID] == nil + && ( + retriesMissingProjections + || !attemptedProjectionTaskIDs.contains(taskID) + ) + ) || ( + retriesMissingProjections + && taskIDs.contains(taskID) + && listSucceeded != true + ) if mustRead { guard reads < 20 else { failed = true break } reads += 1 + attemptedProjectionTaskIDs.insert(taskID) do { if let projection = try await projectionSource.read( threadID: taskID ) { - guard generation == stateGeneration else { + guard generation == stateGeneration, + refreshGeneration == self.refreshGeneration else { return false } projections[taskID] = projection @@ -757,7 +1174,7 @@ actor LocalActivityCollector { } private func taskID(in state: FileState) -> String? { - state.facts.compactMap { fact in + state.normalization.context?.taskID ?? state.facts.compactMap { fact in guard fact.key == .task, fact.availability == .available, case let .identifier(taskID) = fact.value else { @@ -806,6 +1223,64 @@ actor LocalActivityCollector { timestampParser.date(from: value) } + private func clearPublishedContent() { + cachedFacts = nil + cachedFactPaths.removeAll() + cachedFactsIntervalStart = nil + pendingFactRestorePaths.removeAll() + publishedFactPaths.removeAll() + lastProjectionListSucceeded = nil + attemptedProjectionTaskIDs.removeAll() + lastPublishedProjectionIdentities = nil + lastObservationSignature = nil + importContinuationPending = false + } + + private func restartPartialFactRestores() { + for path in Array(files.keys) { + guard var state = files[path], + state.factRestoreFileSize != nil else { + continue + } + state.facts.removeAll(keepingCapacity: false) + state.eventIDs.removeAll(keepingCapacity: false) + state.factsLoaded = false + state.factRestoreOffset = 0 + state.factRestoreFileSize = nil + state.restoredFactIdentities.removeAll(keepingCapacity: false) + files[path] = state + } + } + + private func unloadPersistedFacts(activePaths: Set) { + guard let directory = stateURL else { return } + for path in Array(files.keys) { + guard var state = files[path], + state.factsLoaded, + !changedPaths.contains(path), + pendingFactWrites[path] == nil else { + continue + } + let file = factsURL( + forFingerprint: state.storageFingerprint + ?? stateFileName(for: path), + in: directory + ) + guard FileManager.default.fileExists(atPath: file.path) else { + continue + } + state.facts.removeAll(keepingCapacity: false) + if !activePaths.contains(path) { + state.eventIDs.removeAll(keepingCapacity: false) + } + state.factsLoaded = false + state.factRestoreOffset = 0 + state.factRestoreFileSize = nil + state.restoredFactIdentities.removeAll(keepingCapacity: false) + files[path] = state + } + } + @discardableResult private func persist() -> Bool { guard !changedPaths.isEmpty else { return true } @@ -830,7 +1305,7 @@ actor LocalActivityCollector { } let data = try JSONEncoder().encode( PersistedFile( - version: 6, + version: 7, path: nil, pathFingerprint: storageFingerprint, cursor: state.cursor, @@ -887,12 +1362,12 @@ actor LocalActivityCollector { : nil var legacyPaths = Set() for entry in entries where entry.pathExtension == "json" { - guard let data = try? Data(contentsOf: entry), + guard let data = Self.readMetadata(at: entry), let file = try? JSONDecoder().decode( PersistedFile.self, from: data ), - [4, 5, 6].contains(file.version) else { + [4, 5, 6, 7].contains(file.version) else { restoreWarning = "Saved local activity could not be read" continue } @@ -908,9 +1383,12 @@ actor LocalActivityCollector { storageFingerprint: entry.deletingPathExtension() .lastPathComponent, factsLoaded: false, - requiresContextRebuild: file.version == 4 + requiresContextRebuild: file.version == 4, + factRestoreOffset: 0, + factRestoreFileSize: nil, + restoredFactIdentities: [] ) - if file.version == 6 { + if file.version >= 6 { let fingerprint = entry.deletingPathExtension() .lastPathComponent guard file.path == nil, @@ -1046,68 +1524,226 @@ actor LocalActivityCollector { } private func persistFacts(_ write: FactWrite, to file: URL) throws { - let data = try write.facts.reduce(into: Data()) { result, fact in - result.append(try JSONEncoder().encode(fact)) - result.append(0x0A) - } if write.rewritesFile { - try data.write(to: file, options: .atomic) + let temporary = file.deletingLastPathComponent() + .appendingPathComponent(".\(UUID().uuidString).tmp") + guard FileManager.default.createFile( + atPath: temporary.path, + contents: nil + ) else { + throw CocoaError(.fileWriteUnknown) + } + do { + try writeFacts(write.facts, to: temporary, appending: false) + if FileManager.default.fileExists(atPath: file.path) { + _ = try FileManager.default.replaceItemAt( + file, + withItemAt: temporary + ) + } else { + try FileManager.default.moveItem( + at: temporary, + to: file + ) + } + } catch { + try? FileManager.default.removeItem(at: temporary) + throw error + } return } if !FileManager.default.fileExists(atPath: file.path) { - try Data().write(to: file, options: .atomic) + guard FileManager.default.createFile( + atPath: file.path, + contents: nil + ) else { + throw CocoaError(.fileWriteUnknown) + } } + try writeFacts(write.facts, to: file, appending: true) + } + + private func writeFacts( + _ facts: [LocalActivityFact], + to file: URL, + appending: Bool + ) throws { let handle = try FileHandle(forWritingTo: file) defer { try? handle.close() } - try handle.seekToEnd() - try handle.write(contentsOf: data) + if appending { + try handle.seekToEnd() + } + let encoder = JSONEncoder() + var buffer = Data() + buffer.reserveCapacity(262_144) + for fact in facts { + var data = try encoder.encode(fact) + data.append(0x0A) + if !buffer.isEmpty, buffer.count + data.count > 262_144 { + try handle.write(contentsOf: buffer) + buffer.removeAll(keepingCapacity: true) + } + if data.count > 262_144 { + try handle.write(contentsOf: data) + } else { + buffer.append(data) + } + } + if !buffer.isEmpty { + try handle.write(contentsOf: buffer) + } } - private func restoreFacts(from file: URL) -> [LocalActivityFact]? { - guard let data = try? Data(contentsOf: file) else { return nil } + private func restoreFacts( + from file: URL, + state: inout FileState, + interval: DateInterval, + maximumLines: Int, + maximumBytes: UInt64 + ) -> FactLoadOutcome { + guard let attributes = try? FileManager.default.attributesOfItem( + atPath: file.path + ), + let fileSize = (attributes[.size] as? NSNumber)?.uint64Value, + let handle = try? FileHandle(forReadingFrom: file) else { + return .invalid + } + defer { try? handle.close() } + if state.factRestoreFileSize != fileSize + || state.factRestoreOffset > fileSize { + state.facts.removeAll(keepingCapacity: false) + state.eventIDs.removeAll(keepingCapacity: false) + state.restoredFactIdentities.removeAll(keepingCapacity: false) + state.factRestoreOffset = 0 + state.factRestoreFileSize = fileSize + } let decoder = JSONDecoder() - var facts: [LocalActivityFact] = [] - var seenFacts = Set() - for line in data.split(separator: 0x0A) { - guard let fact = try? decoder.decode( - LocalActivityFact.self, - from: Data(line) - ) else { - return nil - } - if let eventID = fact.eventID { - let identity = "\(eventID)|\(fact.key.rawValue)" - guard seenFacts.insert(identity).inserted else { continue } + var isValid = true + guard let result = try? BoundedJSONLReader.read( + handle: handle, + from: state.factRestoreOffset, + maximumLines: maximumLines, + maximumBytes: maximumBytes, + maximumRecordBytes: Self.maximumMetadataBytes, + discardsPartialRecordAtByteLimit: false, + onLine: { line, _ in + guard isValid else { return false } + guard let fact = try? decoder.decode( + LocalActivityFact.self, + from: line + ) else { + isValid = false + return false + } + guard factPassesHistoryCutoff(fact) else { + return true + } + let isActive = isFactActive(fact, in: interval) + if let eventID = fact.eventID { + guard isActive else { + return true + } + state.eventIDs.insert(eventID) + let identity = FactIdentity( + eventID: eventID, + key: fact.key.rawValue + ) + guard state.restoredFactIdentities.insert( + identity + ).inserted else { + return true + } + } + if isActive { + appendRestoredFact(fact, to: &state.facts) + } + return true } + ), + isValid, + result.oversizedRecordCount == 0 else { + return .invalid + } + state.factRestoreOffset = result.resumeByteOffset + if result.completeByteOffset == fileSize { + state.factRestoreOffset = 0 + state.factRestoreFileSize = nil + state.restoredFactIdentities.removeAll(keepingCapacity: false) + return .ready( + linesRead: result.processedLineCount, + bytesRead: result.bytesRead + ) + } + guard result.stoppedEarly else { return .invalid } + return .partial( + linesRead: result.processedLineCount, + bytesRead: result.bytesRead + ) + } + + private func appendRestoredFact( + _ fact: LocalActivityFact, + to facts: inout [LocalActivityFact] + ) { + guard fact.key == .context, + case let .tokens(usage) = fact.value, + let eventID = fact.eventID, + let lastIndex = facts.indices.last, + facts[lastIndex].key == .token, + facts[lastIndex].eventID == eventID, + facts[lastIndex].eventTimestamp == fact.eventTimestamp, + facts[lastIndex].source == fact.source, + facts[lastIndex].context == fact.context else { facts.append(fact) + return } - return facts + facts[lastIndex].contextUsage = usage } - private func loadFactsIfNeeded(for path: String) { + private func loadFactsIfNeeded( + for path: String, + interval: DateInterval, + maximumLines: Int, + maximumBytes: UInt64 + ) -> FactLoadOutcome { guard var state = files[path], !state.factsLoaded, let directory = stateURL else { - return + return .ready(linesRead: 0, bytesRead: 0) } - guard let restoredFacts = restoreFacts( + let outcome = restoreFacts( from: factsURL( forFingerprint: state.storageFingerprint ?? stateFileName(for: path), in: directory - ) - ) else { - files[path] = nil - return + ), + state: &state, + interval: interval, + maximumLines: maximumLines, + maximumBytes: maximumBytes + ) + guard case .invalid = outcome else { + if case .ready = outcome { + if let pending = pendingFactWrites[path], + !pending.rewritesFile { + for fact in pending.facts { + if let eventID = fact.eventID { + guard state.eventIDs.insert(eventID).inserted + else { continue } + } + if factPassesHistoryCutoff(fact), + isFactActive(fact, in: interval) { + state.facts.append(fact) + } + } + } + state.factsLoaded = true + } + files[path] = state + return outcome } - let facts = factsAfterHistoryCutoff(restoredFacts) - state.facts = facts - state.eventIDs = Set(facts.compactMap(\.eventID)) - let activityBounds = tokenActivityBounds(facts) - state.activityStart = activityBounds?.start - state.activityEnd = activityBounds?.end - state.factsLoaded = true - files[path] = state + files[path] = nil + return .invalid } private func recordsAfterHistoryCutoff( @@ -1129,27 +1765,68 @@ actor LocalActivityCollector { private func factsAfterHistoryCutoff( _ facts: [LocalActivityFact] ) -> [LocalActivityFact] { - guard let historyCutoff else { return facts } - return facts.filter { fact in - guard let timestamp = fact.eventTimestamp, - let date = parseTimestamp(timestamp) else { - return fact.eventID == nil + facts.filter(factPassesHistoryCutoff) + } + + private func factPassesHistoryCutoff( + _ fact: LocalActivityFact + ) -> Bool { + guard let historyCutoff else { return true } + guard let timestamp = fact.eventTimestamp, + let date = parseTimestamp(timestamp) else { + return fact.eventID == nil + } + return date >= historyCutoff + } + + private func isFactActive( + _ fact: LocalActivityFact, + in interval: DateInterval + ) -> Bool { + switch fact.key { + case .task, .parent, .root, .agent: + return true + default: + break + } + if case let .duration(duration) = fact.value { + return duration.completedAt > interval.start + } + if case let .turnTiming(timing) = fact.value { + if let completedAt = timing.completedAt { + return completedAt > interval.start } - return date >= historyCutoff + return timing.startedAt.map { $0 >= interval.start } ?? false + } + guard let timestamp = fact.eventTimestamp, + let date = parseTimestamp(timestamp) else { + return fact.eventID == nil } + return date >= interval.start + } + + private func factsForActiveInterval( + _ facts: [LocalActivityFact], + interval: DateInterval + ) -> [LocalActivityFact] { + facts.filter { isFactActive($0, in: interval) } } private func tokenActivityBounds( _ facts: [LocalActivityFact] ) -> DateInterval? { - let dates = facts.compactMap { fact -> Date? in + var start: Date? + var end: Date? + for fact in facts { guard fact.key == .token, - let timestamp = fact.eventTimestamp else { - return nil + let timestamp = fact.eventTimestamp, + let date = parseTimestamp(timestamp) else { + continue } - return parseTimestamp(timestamp) + start = start.map { min($0, date) } ?? date + end = end.map { max($0, date) } ?? date } - guard let start = dates.min(), let end = dates.max() else { return nil } + guard let start, let end else { return nil } return DateInterval(start: start, end: end) } @@ -1160,7 +1837,7 @@ actor LocalActivityCollector { FileManager.default.fileExists(atPath: file.path) else { return .missing } - guard let data = try? Data(contentsOf: file), + guard let data = readMetadata(at: file), let marker = try? JSONDecoder().decode( DeletionMarker.self, from: data @@ -1177,10 +1854,35 @@ actor LocalActivityCollector { let file = stateDirectory.appendingPathComponent( "deletion-cutoff.json" ) - guard let data = try? Data(contentsOf: file) else { return nil } + guard let data = readMetadata(at: file) else { return nil } return try? JSONDecoder().decode(Date.self, from: data) } + private static func readMetadata(at file: URL) -> Data? { + guard let handle = try? FileHandle(forReadingFrom: file) else { + return nil + } + defer { try? handle.close() } + var data = Data() + while data.count <= maximumMetadataBytes { + let remaining = maximumMetadataBytes + 1 - data.count + let chunk: Data + do { + guard let value = try handle.read( + upToCount: min(65_536, remaining) + ) else { + return data + } + chunk = value + } catch { + return nil + } + guard !chunk.isEmpty else { return data } + data.append(chunk) + } + return nil + } + private static func deletionMarkerURL( for stateDirectory: URL ) -> URL? { diff --git a/Sources/CodexLimits/LocalActivityFactIndex.swift b/Sources/CodexLimits/LocalActivityFactIndex.swift index 0b067b2..3b29871 100644 --- a/Sources/CodexLimits/LocalActivityFactIndex.swift +++ b/Sources/CodexLimits/LocalActivityFactIndex.swift @@ -1,40 +1,40 @@ import Foundation struct LocalActivityFactIndex { - private let entries: [(date: Date, fact: LocalActivityFact)] + private let entries: [(date: Date, index: Int)] private let turnEntries: [ - (start: Date, end: Date, fact: LocalActivityFact) + (start: Date, end: Date, index: Int) ] private let boundaryEntries: [ ( date: Date, affectedStart: Date?, affectedEnd: Date?, - fact: LocalActivityFact + index: Int ) ] init(_ facts: [LocalActivityFact]) { let parser = LocalEventTimestampParser() - entries = facts.compactMap { fact in + entries = facts.enumerated().compactMap { index, fact in guard let timestamp = fact.eventTimestamp, let date = parser.date(from: timestamp) else { return nil } - return (date, fact) + return (date, index) } .sorted { $0.date < $1.date } - turnEntries = facts.compactMap { fact in + turnEntries = facts.enumerated().compactMap { index, fact in guard case let .turnTiming(timing) = fact.value, let start = timing.startedAt, let end = timing.completedAt, start < end else { return nil } - return (start, end, fact) + return (start, end, index) } .sorted { $0.start < $1.start } - boundaryEntries = facts.compactMap { fact in + boundaryEntries = facts.enumerated().compactMap { index, fact in guard case let .turnTiming(timing) = fact.value else { return nil } @@ -52,21 +52,29 @@ struct LocalActivityFactIndex { date, min(start, end), max(start, end), - fact + index ) } - return (date, start, end, fact) + return (date, start, end, index) } } - func facts(in interval: DateInterval) -> [LocalActivityFact] { + func facts( + in interval: DateInterval, + from facts: [LocalActivityFact] + ) -> [LocalActivityFact] { let start = lowerBound(for: interval.start) let end = lowerBound(for: interval.end) - return entries[start ..< end].map(\.fact) + return entries[start ..< end].compactMap { + facts.indices.contains($0.index) ? facts[$0.index] : nil + } } - func activityFacts(in interval: DateInterval) -> [LocalActivityFact] { - var result = facts(in: interval).filter { fact in + func activityFacts( + in interval: DateInterval, + from facts: [LocalActivityFact] + ) -> [LocalActivityFact] { + var result = self.facts(in: interval, from: facts).filter { fact in if case .turnTiming = fact.value { return false } return true } @@ -74,7 +82,9 @@ struct LocalActivityFactIndex { result += turnEntries[.. interval.start } - .map(\.fact) + .compactMap { + facts.indices.contains($0.index) ? facts[$0.index] : nil + } result += boundaryEntries .lazy .filter { entry in @@ -89,7 +99,9 @@ struct LocalActivityFactIndex { return interval.contains(entry.date) } } - .map(\.fact) + .compactMap { + facts.indices.contains($0.index) ? facts[$0.index] : nil + } return result } diff --git a/Sources/CodexLimits/LocalActivityNormalizer.swift b/Sources/CodexLimits/LocalActivityNormalizer.swift index 0c181ca..df4aa2f 100644 --- a/Sources/CodexLimits/LocalActivityNormalizer.swift +++ b/Sources/CodexLimits/LocalActivityNormalizer.swift @@ -49,7 +49,7 @@ enum LocalActivityFactValue: Codable, Equatable, Sendable { case duration(LocalActivityDuration) } -struct LocalAgentIdentity: Codable, Equatable, Sendable { +struct LocalAgentIdentity: Codable, Equatable, Hashable, Sendable { let nickname: String? let role: String? } @@ -66,7 +66,7 @@ struct LocalActivityDuration: Codable, Equatable, Sendable { let completedAt: Date } -struct LocalActivityContext: Codable, Equatable, Sendable { +struct LocalActivityContext: Codable, Equatable, Hashable, Sendable { let taskID: String? let turnID: String? let agent: LocalAgentIdentity? @@ -87,6 +87,7 @@ struct LocalActivityFact: Codable, Equatable, Sendable { let source: LocalActivitySourceMetadata var context: LocalActivityContext? = nil var tokenDelta: LocalTokenUsage? = nil + var contextUsage: LocalTokenUsage? = nil } struct LocalActivityNormalizationState: Codable, Equatable, Sendable { @@ -138,7 +139,9 @@ struct LocalActivityNormalizer { if sourceChanged { state.lastTotalTokens = nil state.lastTokenUsage = nil - state.tokenSegment += 1 + if state.tokenSegment < .max { + state.tokenSegment += 1 + } state.context = nil } state.sourceGeneration = sourceGeneration @@ -284,14 +287,18 @@ struct LocalActivityNormalizer { if let tokenUsage = record.tokenUsage { let totalTokens = tokenUsage.totalTokens let delta = state.lastTotalTokens.flatMap { previous in - totalTokens >= previous ? totalTokens - previous : nil + totalTokens >= previous + ? difference(totalTokens, previous) + : nil } let reason: String? if state.lastTotalTokens == nil { reason = sourceChanged ? "source-discontinuity" : "segment-baseline" } else if delta == nil { reason = "cumulative-counter-decreased" - state.tokenSegment += 1 + if state.tokenSegment < .max { + state.tokenSegment += 1 + } } else { reason = nil } @@ -309,13 +316,15 @@ struct LocalActivityNormalizer { context: context, tokenDelta: state.lastTokenUsage.flatMap { tokenDelta(from: $0, to: tokenUsage) - } + }, + contextUsage: record.contextTokenUsage ) ) state.lastTotalTokens = totalTokens state.lastTokenUsage = tokenUsage } - if let contextTokenUsage = record.contextTokenUsage { + if record.tokenUsage == nil, + let contextTokenUsage = record.contextTokenUsage { facts.append( LocalActivityFact( key: .context, @@ -412,10 +421,8 @@ struct LocalActivityNormalizer { from previous: LocalTokenUsage, to current: LocalTokenUsage ) -> LocalTokenUsage? { - let observed = previous.observedComponents.intersection( - current.observedComponents - ) - guard observed.contains(.total), + let observedMask = previous.sharedObservedComponentMask(with: current) + guard previous.observes(.total), current.observes(.total), let total = difference( current.totalTokens, previous.totalTokens @@ -427,7 +434,9 @@ struct LocalActivityNormalizer { _ currentValue: Int64, _ previousValue: Int64 ) -> Int64? { - guard observed.contains(key) else { return 0 } + guard previous.observes(key), current.observes(key) else { + return 0 + } return difference(currentValue, previousValue) } guard @@ -466,7 +475,7 @@ struct LocalActivityNormalizer { outputTokens: output, reasoningOutputTokens: reasoning, totalTokens: total, - observedComponents: observed + observedComponentMask: observedMask ) } diff --git a/Sources/CodexLimits/LocalTokenActivity.swift b/Sources/CodexLimits/LocalTokenActivity.swift index 5db4ca7..afc8515 100644 --- a/Sources/CodexLimits/LocalTokenActivity.swift +++ b/Sources/CodexLimits/LocalTokenActivity.swift @@ -51,7 +51,7 @@ struct LocalTokenActivitySnapshot: Equatable, Sendable { _ reason: String, interval: DateInterval ) -> LocalTokenActivitySnapshot { - LocalTokenActivitySnapshot( + return LocalTokenActivitySnapshot( tokens: nil, interval: interval, coverage: .unavailable, @@ -63,6 +63,42 @@ struct LocalTokenActivitySnapshot: Equatable, Sendable { ) } + func updating( + interval: DateInterval, + observation: LocalActivityObservation + ) -> LocalTokenActivitySnapshot { + let source: (version: String?, observedAt: Date?) = switch observation { + case let .continuous(version, observedAt), + let .gap(version, observedAt, _): + (version, observedAt) + case .unavailable: + (nil, nil) + } + let updatedCoverage: CoverageLevel + let updatedReason: String? + switch observation { + case .continuous: + updatedCoverage = coverage + updatedReason = reason + case let .gap(_, _, reason): + updatedCoverage = .low + updatedReason = reason + case let .unavailable(reason): + updatedCoverage = .unavailable + updatedReason = reason + } + return LocalTokenActivitySnapshot( + tokens: tokens, + interval: interval, + coverage: updatedCoverage, + reason: updatedReason, + sourceVersion: source.version, + observedAt: source.observedAt, + points: points, + accountComparison: accountComparison + ) + } + func slice(in selectedInterval: DateInterval) -> LocalTokenActivitySlice { let baseline = points.last { $0.date < selectedInterval.start diff --git a/Sources/CodexLimits/LocalWorkloadMix.swift b/Sources/CodexLimits/LocalWorkloadMix.swift index a7e15b3..e8e5fa2 100644 --- a/Sources/CodexLimits/LocalWorkloadMix.swift +++ b/Sources/CodexLimits/LocalWorkloadMix.swift @@ -1,20 +1,13 @@ import Foundation struct LocalEventTimestampParser { - private let fractional: ISO8601DateFormatter - private let standard: ISO8601DateFormatter - - init() { - fractional = ISO8601DateFormatter() - fractional.formatOptions = [ - .withInternetDateTime, - .withFractionalSeconds - ] - standard = ISO8601DateFormatter() - } + private let fractional = Date.ISO8601FormatStyle( + includingFractionalSeconds: true + ) + private let standard = Date.ISO8601FormatStyle() func date(from value: String) -> Date? { - fractional.date(from: value) ?? standard.date(from: value) + (try? fractional.parse(value)) ?? (try? standard.parse(value)) } } diff --git a/Sources/CodexLimits/MenuContentView.swift b/Sources/CodexLimits/MenuContentView.swift index 9868520..203a8c9 100644 --- a/Sources/CodexLimits/MenuContentView.swift +++ b/Sources/CodexLimits/MenuContentView.swift @@ -75,7 +75,7 @@ struct MenuContentView: View { .padding(.vertical, 12) } .frame(width: layout.width, height: layout.height) - .task { await monitor.refresh() } + .task { await monitor.refresh(forceHistorySync: false) } .environment(\.locale, Locale(identifier: "en_US")) } @@ -1173,6 +1173,10 @@ private struct ConcurrencyWorkspace: View { } var body: some View { + content(slice) + } + + private func content(_ slice: ActivityTimelineSlice) -> some View { VStack(alignment: .leading, spacing: 14) { VStack(alignment: .leading, spacing: 4) { Text("Concurrency") @@ -1211,8 +1215,8 @@ private struct ConcurrencyWorkspace: View { } .frame(minHeight: 190) } else { - chart - selectedPointDetail + chart(slice) + selectedPointDetail(slice) zoomControls } @@ -1260,7 +1264,7 @@ private struct ConcurrencyWorkspace: View { .foregroundStyle(.tertiary) } - private var chart: some View { + private func chart(_ slice: ActivityTimelineSlice) -> some View { Chart { ForEach(slice.points) { point in LineMark( @@ -1309,7 +1313,8 @@ private struct ConcurrencyWorkspace: View { selectNearestPoint( at: location, proxy: proxy, - geometry: geometry + geometry: geometry, + points: slice.points ) case .ended: selectedPoint = nil @@ -1320,15 +1325,15 @@ private struct ConcurrencyWorkspace: View { .frame(height: 260) .accessibilityElement(children: .ignore) .accessibilityLabel("Concurrency") - .accessibilityValue(accessibilityValue) + .accessibilityValue(accessibilityValue(slice)) .accessibilityHint( "Use Previous point and Next point for exact values." ) } - private var accessibilityValue: String { + private func accessibilityValue(_ slice: ActivityTimelineSlice) -> String { if let selectedPoint { - return pointSummary(selectedPoint) + return pointSummary(selectedPoint, coverage: slice.coverage) } let trees = slice.maximumConcurrency == 1 ? "Active Task Tree" @@ -1336,7 +1341,9 @@ private struct ConcurrencyWorkspace: View { return "\(duration(slice.activeTime)) Active Time, peak \(slice.maximumConcurrency) \(trees), \(slice.coverage.displayName) coverage" } - private var selectedPointDetail: some View { + private func selectedPointDetail( + _ slice: ActivityTimelineSlice + ) -> some View { HStack(spacing: 10) { if let selectedPoint { let taskTrees = taskTreeLabels(selectedPoint) @@ -1371,13 +1378,13 @@ private struct ConcurrencyWorkspace: View { } Spacer() Button { - moveSelection(by: -1) + moveSelection(in: slice.points, by: -1) } label: { Image(systemName: "chevron.left") } .accessibilityLabel("Previous point") Button { - moveSelection(by: 1) + moveSelection(in: slice.points, by: 1) } label: { Image(systemName: "chevron.right") } @@ -1433,18 +1440,24 @@ private struct ConcurrencyWorkspace: View { } } - private func pointSummary(_ point: ConcurrencyPoint) -> String { + private func pointSummary( + _ point: ConcurrencyPoint, + coverage: CoverageLevel + ) -> String { let trees = taskTreeLabels(point).joined(separator: ", ") let countLabel = point.count == 1 ? "Active Task Tree" : "Active Task Trees" let treeDetail = trees.isEmpty ? "No Active Task Trees" : trees - return "\(point.count) \(countLabel) at \(point.date.formatted(date: .abbreviated, time: .shortened)). \(treeDetail). Codex local records. \(slice.coverage.displayName) coverage." + return "\(point.count) \(countLabel) at \(point.date.formatted(date: .abbreviated, time: .shortened)). \(treeDetail). Codex local records. \(coverage.displayName) coverage." } - private func moveSelection(by offset: Int) { + private func moveSelection( + in points: [ConcurrencyPoint], + by offset: Int + ) { selectedPoint = steppedPoint( - in: slice.points, + in: points, from: selectedPoint, by: offset ) @@ -1453,7 +1466,8 @@ private struct ConcurrencyWorkspace: View { private func selectNearestPoint( at location: CGPoint, proxy: ChartProxy, - geometry: GeometryProxy + geometry: GeometryProxy, + points: [ConcurrencyPoint] ) { guard let date = chartDate( at: location, @@ -1461,7 +1475,7 @@ private struct ConcurrencyWorkspace: View { geometry: geometry ) else { return } selectedPoint = nearestPoint( - in: slice.points, + in: points, to: date, date: \.date ) @@ -1517,16 +1531,10 @@ private struct TokenActivityWorkspace: View { guard !store.state.filters.isEmpty else { return reader.localTokenActivity.slice(in: visibleRange) } - let receipts = reader.usageReceipts.slice( + return reader.usageReceipts.localTokenSlice( in: visibleRange, filters: store.state.filters ) - return LocalTokenActivitySlice( - tokens: receipts.totalTokens, - points: receipts.points, - coverage: receipts.coverage, - reason: receipts.reason - ) } private var accountCoversVisibleRange: Bool { @@ -1538,6 +1546,10 @@ private struct TokenActivityWorkspace: View { } var body: some View { + content(localSlice) + } + + private func content(_ slice: LocalTokenActivitySlice) -> some View { VStack(alignment: .leading, spacing: 16) { VStack(alignment: .leading, spacing: 4) { Text("Token activity") @@ -1552,11 +1564,11 @@ private struct TokenActivityWorkspace: View { ViewThatFits(in: .horizontal) { HStack(alignment: .top, spacing: 12) { accountCard - localCard + localCard(slice) } VStack(spacing: 12) { accountCard - localCard + localCard(slice) } } @@ -1573,20 +1585,20 @@ private struct TokenActivityWorkspace: View { } } - if localSlice.points.isEmpty { + if slice.points.isEmpty { WorkspaceMessage( icon: "chart.xyaxis.line", title: "No local token activity", - message: localEmptyMessage + message: localEmptyMessage(slice) ) { EmptyView() } .frame(minHeight: 170) } else { - localChart + localChart(slice) } - selectedPointDetail + selectedPointDetail(slice) } } .onChange(of: visibleRange) { _, range in @@ -1622,15 +1634,15 @@ private struct TokenActivityWorkspace: View { ) } - private var localCard: some View { + private func localCard(_ slice: LocalTokenActivitySlice) -> some View { TokenSourceCard( title: "Local", source: "Local Codex records", value: reader.localTokenActivity.tokens == nil ? "Not available" - : compactTokenCount(localSlice.tokens), - detail: localDetail, - coverage: coverageName(localSlice.coverage), + : compactTokenCount(slice.tokens), + detail: localDetail(slice), + coverage: coverageName(slice.coverage), freshness: reader.localTokenActivity.observedAt, freshnessLabel: "Updated", color: .purple @@ -1686,12 +1698,12 @@ private struct TokenActivityWorkspace: View { } } - private var localDetail: String { + private func localDetail(_ slice: LocalTokenActivitySlice) -> String { var details: [String] = [] if let version = reader.localTokenActivity.sourceVersion { details.append("Codex \(version)") } - if let reason = localSlice.reason { + if let reason = slice.reason { details.append(readerFacingLocalReason(reason)) } return details.isEmpty @@ -1699,22 +1711,25 @@ private struct TokenActivityWorkspace: View { : details.joined(separator: " · ") } - private var localEmptyMessage: String { - localSlice.reason.map(readerFacingLocalReason) + private func localEmptyMessage(_ slice: LocalTokenActivitySlice) -> String { + slice.reason.map(readerFacingLocalReason) ?? "No local token events were found in this range." } - private var chartPoints: [LocalTokenActivityPoint] { - if localSlice.points.first?.date == visibleRange.start { - return localSlice.points + private func renderedChartPoints( + _ slice: LocalTokenActivitySlice + ) -> [LocalTokenActivityPoint] { + let points = slice.points + if points.first?.date == visibleRange.start { + return downsampledForDisplay(points) } return [LocalTokenActivityPoint(date: visibleRange.start, tokens: 0)] - + localSlice.points + + downsampledForDisplay(points, limit: 999) } - private var localChart: some View { + private func localChart(_ slice: LocalTokenActivitySlice) -> some View { Chart { - ForEach(chartPoints) { point in + ForEach(renderedChartPoints(slice)) { point in LineMark( x: .value("Time", point.date), y: .value("Local tokens", point.tokens), @@ -1761,7 +1776,8 @@ private struct TokenActivityWorkspace: View { selectNearestPoint( at: location, proxy: proxy, - geometry: geometry + geometry: geometry, + points: slice.points ) case .ended: selectedPoint = nil @@ -1775,7 +1791,7 @@ private struct TokenActivityWorkspace: View { .accessibilityValue( selectedPoint.map { "\(compactTokenCount($0.tokens)) local tokens, \($0.date.formatted(date: .abbreviated, time: .shortened))" - } ?? "\(compactTokenCount(localSlice.tokens)) local tokens in the selected range" + } ?? "\(compactTokenCount(slice.tokens)) local tokens in the selected range" ) .accessibilityHint( "Use Previous point and Next point for exact values." @@ -1783,7 +1799,9 @@ private struct TokenActivityWorkspace: View { } @ViewBuilder - private var selectedPointDetail: some View { + private func selectedPointDetail( + _ slice: LocalTokenActivitySlice + ) -> some View { VStack(alignment: .leading, spacing: 7) { HStack(spacing: 10) { if let selectedPoint { @@ -1823,13 +1841,13 @@ private struct TokenActivityWorkspace: View { } Spacer() Button { - moveSelection(by: -1) + moveSelection(in: slice.points, by: -1) } label: { Image(systemName: "chevron.left") } .accessibilityLabel("Previous point") Button { - moveSelection(by: 1) + moveSelection(in: slice.points, by: 1) } label: { Image(systemName: "chevron.right") } @@ -1855,9 +1873,12 @@ private struct TokenActivityWorkspace: View { ) } - private func moveSelection(by offset: Int) { + private func moveSelection( + in points: [LocalTokenActivityPoint], + by offset: Int + ) { selectedPoint = steppedPoint( - in: localSlice.points, + in: points, from: selectedPoint, by: offset ) @@ -1885,8 +1906,14 @@ private struct TokenActivityWorkspace: View { "Saved local activity could not be read" case "Local activity could not be saved": "Local activity could not be saved" + case "Local task import is still in progress": + "Local activity is still loading" case "Local task record continuity changed": "A local task record changed" + case "Local activity read was cancelled": + "Local activity could not finish loading" + case "Account changed during local activity read": + "Local activity changed while loading" default: reason } @@ -1895,7 +1922,8 @@ private struct TokenActivityWorkspace: View { private func selectNearestPoint( at location: CGPoint, proxy: ChartProxy, - geometry: GeometryProxy + geometry: GeometryProxy, + points: [LocalTokenActivityPoint] ) { guard let date = chartDate( at: location, @@ -1903,7 +1931,7 @@ private struct TokenActivityWorkspace: View { geometry: geometry ) else { return } selectedPoint = nearestPoint( - in: localSlice.points, + in: points, to: date, date: \.date ) @@ -2316,7 +2344,11 @@ private struct UsageRemainingChart: View { @ChartContentBuilder private var observedMarks: some ChartContent { ForEach( - Array(chart.observedSegments(within: visibleRange).enumerated()), + Array( + chart.observedSegments(within: visibleRange) + .map { downsampledForDisplay($0) } + .enumerated() + ), id: \.offset ) { segmentIndex, segment in ForEach(segment) { point in @@ -2770,11 +2802,13 @@ private struct FactsWorkspace: View { } private var activeTimeContent: some View { - let slice = reader.activityTimeline.slice( - in: reader.activityTimeline.interval, - filters: store.state.filters - ) let availability = reader.activeTimeAvailability + let filteredSlice = store.state.filters.isEmpty + ? nil + : reader.activityTimeline.slice( + in: reader.activityTimeline.interval, + filters: store.state.filters + ) return VStack(alignment: .leading, spacing: 9) { FactRow( label: "Active time this week", @@ -2806,20 +2840,23 @@ private struct FactsWorkspace: View { } FactRow( label: "Peak concurrency", - value: "\(slice.maximumConcurrency)" + value: "\(filteredSlice?.maximumConcurrency ?? availability.maximumConcurrency)" ) FactRow( label: "Waiting", - value: slice.waitingTime.map { + value: (filteredSlice?.waitingTime + ?? availability.waitingTime).map { duration(Int64($0.rounded(.down))) } ?? "Unavailable" ) FactRow( label: "Polling", - value: slice.pollingTime.map { + value: (filteredSlice?.pollingTime + ?? availability.pollingTime).map { duration(Int64($0.rounded(.down))) } ?? "Unavailable", - detail: slice.activityBreakdownReason + detail: filteredSlice?.activityBreakdownReason + ?? availability.activityBreakdownReason ) FactRow( label: "Source", @@ -2902,44 +2939,54 @@ private struct FactsWorkspace: View { @ViewBuilder private var receiptContent: some View { - let slice = reader.usageReceipts.slice( + let overview = reader.usageReceipts.overview( in: receiptRange, filters: store.state.filters ) - if slice.receipts.isEmpty { - Text(slice.reason ?? "No Usage Receipts are available.") + if overview.receipts.isEmpty { + Text(overview.reason ?? "No Usage Receipts are available.") .foregroundStyle(.secondary) } else { HStack { - Text("\(compactTokenCount(slice.totalTokens)) local tokens") + Text( + "\(compactTokenCount(overview.totalTokens)) local tokens" + ) Spacer() - Text("\(slice.receiptCoverage.displayName) coverage") + Text("\(overview.receiptCoverage.displayName) coverage") .foregroundStyle(.secondary) } .font(.caption) .accessibilityElement(children: .combine) - if let receiptReason = slice.receiptReason { + if let receiptReason = overview.receiptReason { Text(receiptReason) .font(.caption) .foregroundStyle(.tertiary) } - if slice.unattributedTokens > 0 { + if overview.unattributedTokens > 0 { Text( - "\(compactTokenCount(slice.unattributedTokens)) local tokens could not be matched to a Task." + "\(compactTokenCount(overview.unattributedTokens)) local tokens could not be matched to a Task." ) .font(.caption) .foregroundStyle(.secondary) } - ForEach(receiptProjects(in: slice), id: \.self) { project in + ForEach(receiptProjects(in: overview), id: \.self) { project in DisclosureGroup { ForEach( - slice.receipts.filter { + overview.receipts.filter { $0.projectLabel == project } - ) { receipt in - receiptDisclosure(receipt) + ) { summary in + UsageReceiptSummaryDisclosure( + summary: summary, + snapshot: reader.usageReceipts, + selectedInterval: receiptRange, + filters: store.state.filters, + contentRevision: + reader.reusableLocalAggregates?.contentRevision + ?? 0 + ) } } label: { Label( @@ -2964,75 +3011,13 @@ private struct FactsWorkspace: View { } private func receiptProjects( - in slice: UsageReceiptSlice + in overview: UsageReceiptOverview ) -> [String?] { - Array(Set(slice.receipts.map(\.projectLabel))).sorted { + Array(Set(overview.receipts.map(\.projectLabel))).sorted { ($0 ?? "") < ($1 ?? "") } } - @ViewBuilder - private func receiptDisclosure(_ receipt: UsageReceipt) -> some View { - DisclosureGroup { - VStack(alignment: .leading, spacing: 8) { - FactRow( - label: "Local Token Activity", - value: compactTokenCount(receipt.tokens) - ) - FactRow( - label: "Task Tree", - value: "\(receipt.taskCount) \(receipt.taskCount == 1 ? "Task" : "Tasks")" - ) - FactRow( - label: "Range", - value: receipt.intervalText - ) - UsageReceiptDiagnosticsView( - diagnostics: receipt.diagnostics - ) - Divider() - VStack(alignment: .leading, spacing: 8) { - Text("Task Tree") - .font(.caption.weight(.semibold)) - .foregroundStyle(.secondary) - Text( - "Open a Task to see its agents and turns. Models are the effective settings recorded for each turn." - ) - .font(.caption) - .foregroundStyle(.tertiary) - UsageReceiptTaskTreeView( - node: receipt.taskTree, - isRoot: true - ) - } - UsageReceiptBreakdownView( - title: "Models", - values: receipt.models - ) - UsageReceiptBreakdownView( - title: "Reasoning", - values: receipt.reasoningLevels - ) - FactRow( - label: "Coverage", - value: receipt.coverage.displayName, - detail: receipt.reason - ) - } - .padding(.leading, 4) - } label: { - HStack { - Text("Task \(receipt.displayTaskID)") - Spacer() - Text(compactTokenCount(receipt.tokens)) - .foregroundStyle(.secondary) - .monospacedDigit() - } - .accessibilityElement(children: .ignore) - .accessibilityLabel(receipt.accessibilityValue) - } - } - @ViewBuilder private func accountFacts(_ facts: AccountFacts) -> some View { if let value = facts.lifetimeTokens { @@ -3125,6 +3110,123 @@ private struct FactsWorkspace: View { } } +private struct UsageReceiptSummaryDisclosure: View { + let summary: UsageReceiptSummary + let snapshot: UsageReceiptSnapshot + let selectedInterval: DateInterval + let filters: WorkspaceFilters + let contentRevision: UInt64 + + @State private var isExpanded = false + @State private var receipt: UsageReceipt? + + var body: some View { + DisclosureGroup(isExpanded: expansion) { + if let receipt { + receiptDetails(receipt) + } else { + Text("Task details are not available.") + .font(.caption) + .foregroundStyle(.secondary) + } + } label: { + HStack { + Text("Task \(summary.displayTaskID)") + Spacer() + Text(compactTokenCount(summary.tokens)) + .foregroundStyle(.secondary) + .monospacedDigit() + } + .accessibilityElement(children: .ignore) + .accessibilityLabel(summary.accessibilityValue) + } + .onChange(of: summary) { _, _ in + reloadIfExpanded() + } + .onChange(of: selectedInterval) { _, _ in + reloadIfExpanded() + } + .onChange(of: filters) { _, _ in + reloadIfExpanded() + } + .onChange(of: contentRevision) { _, _ in + reloadIfExpanded() + } + } + + private var expansion: Binding { + Binding( + get: { isExpanded }, + set: { expanded in + isExpanded = expanded + receipt = expanded ? loadReceipt() : nil + } + ) + } + + private func loadReceipt() -> UsageReceipt? { + snapshot.receipt( + rootTaskID: summary.rootTaskID, + in: selectedInterval, + filters: filters + ) + } + + private func reloadIfExpanded() { + guard isExpanded else { return } + receipt = loadReceipt() + } + + private func receiptDetails(_ receipt: UsageReceipt) -> some View { + VStack(alignment: .leading, spacing: 8) { + FactRow( + label: "Local Token Activity", + value: compactTokenCount(receipt.tokens) + ) + FactRow( + label: "Task Tree", + value: "\(receipt.taskCount) \(receipt.taskCount == 1 ? "Task" : "Tasks")" + ) + FactRow( + label: "Range", + value: receipt.intervalText + ) + UsageReceiptDiagnosticsView( + diagnostics: receipt.diagnostics + ) + Divider() + VStack(alignment: .leading, spacing: 8) { + Text("Task Tree") + .font(.caption.weight(.semibold)) + .foregroundStyle(.secondary) + Text( + "Open a Task to see its agents and turns. Models are the effective settings recorded for each turn." + ) + .font(.caption) + .foregroundStyle(.tertiary) + UsageReceiptTaskTreeView( + node: receipt.taskTree, + isRoot: true + ) + } + UsageReceiptBreakdownView( + title: "Models", + values: receipt.models + ) + UsageReceiptBreakdownView( + title: "Reasoning", + values: receipt.reasoningLevels + ) + FactRow( + label: "Coverage", + value: receipt.coverage.displayName, + detail: receipt.reason + ) + } + .padding(.leading, 4) + } +} + private func compactTokenCount(_ value: Int64) -> String { value.formatted( .number.notation(.compactName).precision(.fractionLength(0 ... 1)) @@ -3141,6 +3243,18 @@ func usageReceiptTokenTotalDetail(_ reconciles: Bool?) -> String { return "Input or output is unavailable" } +func downsampledForDisplay( + _ points: [Point], + limit requestedLimit: Int = 1_000 +) -> [Point] { + let limit = max(requestedLimit, 2) + guard points.count > limit else { return points } + let step = Double(points.count - 1) / Double(limit - 1) + return (0 ..< limit).map { + points[min(Int((Double($0) * step).rounded()), points.count - 1)] + } +} + private func steppedPoint( in points: [Point], from selected: Point?, @@ -3153,15 +3267,42 @@ private func steppedPoint( return points[min(max(index + offset, 0), points.count - 1)] } -private func nearestPoint( +func nearestPoint( in points: [Point], to target: Date, - date: KeyPath + date: KeyPath, + within interval: DateInterval? = nil ) -> Point? { - points.min { - abs($0[keyPath: date].timeIntervalSince(target)) - < abs($1[keyPath: date].timeIntervalSince(target)) - } + func lowerBound(_ target: Date, from start: Int, to end: Int) -> Int { + var lower = start + var upper = end + while lower < upper { + let middle = lower + (upper - lower) / 2 + if points[middle][keyPath: date] < target { + lower = middle + 1 + } else { + upper = middle + } + } + return lower + } + let start = interval.map { + lowerBound($0.start, from: 0, to: points.count) + } ?? 0 + let end = interval.map { + let index = lowerBound($0.end, from: start, to: points.count) + return index < points.count + && points[index][keyPath: date] == $0.end ? index + 1 : index + } ?? points.count + guard start < end else { return nil } + let insertion = lowerBound(target, from: start, to: end) + let candidates = [insertion - 1, insertion].filter { + $0 >= start && $0 < end + } + return candidates.min { + abs(points[$0][keyPath: date].timeIntervalSince(target)) + < abs(points[$1][keyPath: date].timeIntervalSince(target)) + }.map { points[$0] } } private func chartDate( diff --git a/Sources/CodexLimits/RolloutTailSource.swift b/Sources/CodexLimits/RolloutTailSource.swift index e2a127f..4e23575 100644 --- a/Sources/CodexLimits/RolloutTailSource.swift +++ b/Sources/CodexLimits/RolloutTailSource.swift @@ -2,6 +2,224 @@ import Foundation enum RolloutTailSourceError: Error, Equatable { case missingFileMetadata + case counterOverflow +} + +struct BoundedJSONLReadResult { + let completeByteOffset: UInt64 + let resumeByteOffset: UInt64 + let bytesRead: UInt64 + let oversizedRecordCount: Int + let processedLineCount: Int + let stoppedEarly: Bool + let discardingOversizedRecord: Bool +} + +struct BoundedJSONLReader { + static let maximumRecordBytes = 16 * 1_024 * 1_024 + private static let chunkBytes = 65_536 + + static func read( + handle: FileHandle, + from startOffset: UInt64, + through endOffset: UInt64? = nil, + maximumLines: Int = .max, + maximumBytes: UInt64 = .max, + maximumRecordBytes: Int = Self.maximumRecordBytes, + startsByDiscardingOversizedRecord: Bool = false, + discardsPartialRecordAtByteLimit: Bool = true, + onLine: (Data, UInt64) throws -> Bool + ) throws -> BoundedJSONLReadResult { + let lineLimit = max(maximumLines, 1) + let byteLimit = max(maximumBytes, 1) + try handle.seek(toOffset: startOffset) + var line = Data() + var lineStart = startOffset + var completeByteOffset = startOffset + var bytesRead: UInt64 = 0 + var oversizedRecordCount = + startsByDiscardingOversizedRecord ? 1 : 0 + var processedLineCount = 0 + var processedBytes: UInt64 = 0 + var skipsCurrentLine = startsByDiscardingOversizedRecord + var countedCurrentOversizedLine = startsByDiscardingOversizedRecord + + while true { + let (currentOffset, currentOffsetOverflowed) = + startOffset.addingReportingOverflow(bytesRead) + guard !currentOffsetOverflowed else { + throw RolloutTailSourceError.counterOverflow + } + let remaining = endOffset.map { + $0 >= currentOffset ? $0 - currentOffset : 0 + } + if remaining == 0 { break } + let budgetRemaining = byteLimit > bytesRead + ? byteLimit - bytesRead + : 0 + if budgetRemaining == 0 { + let discardsPartial = discardsPartialRecordAtByteLimit + && !line.isEmpty + if discardsPartial, !countedCurrentOversizedLine { + oversizedRecordCount += 1 + } + return BoundedJSONLReadResult( + completeByteOffset: completeByteOffset, + resumeByteOffset: discardsPartial || skipsCurrentLine + ? currentOffset + : completeByteOffset, + bytesRead: bytesRead, + oversizedRecordCount: oversizedRecordCount, + processedLineCount: processedLineCount, + stoppedEarly: true, + discardingOversizedRecord: skipsCurrentLine + || discardsPartial + ) + } + let count = min( + Self.chunkBytes, + remaining.map { Int(min($0, UInt64(Self.chunkBytes))) } + ?? Self.chunkBytes, + Int(min(budgetRemaining, UInt64(Self.chunkBytes))) + ) + guard count > 0, + let chunk = try handle.read(upToCount: count), + !chunk.isEmpty else { + break + } + let chunkStart = currentOffset + let (nextBytesRead, overflowed) = bytesRead.addingReportingOverflow( + UInt64(chunk.count) + ) + guard !overflowed else { + throw RolloutTailSourceError.counterOverflow + } + bytesRead = nextBytesRead + var pieceStart = chunk.startIndex + while let newline = chunk[pieceStart...].firstIndex(of: 0x0A) { + let piece = chunk[pieceStart.. maximumRecordBytes { + line.removeAll(keepingCapacity: false) + skipsCurrentLine = true + if !countedCurrentOversizedLine { + oversizedRecordCount += 1 + countedCurrentOversizedLine = true + } + } else { + line.append(contentsOf: piece) + } + } + let distance = chunk.distance( + from: chunk.startIndex, + to: newline + ) + 1 + let (lineEnd, offsetOverflowed) = + chunkStart.addingReportingOverflow(UInt64(distance)) + guard !offsetOverflowed else { + throw RolloutTailSourceError.counterOverflow + } + let shouldContinue: Bool + if skipsCurrentLine { + shouldContinue = true + } else if !line.isEmpty { + shouldContinue = try onLine(line, lineStart) + } else { + shouldContinue = true + } + let (lineBytes, lineBytesOverflowed) = + lineEnd.subtractingReportingOverflow(lineStart) + guard !lineBytesOverflowed else { + throw RolloutTailSourceError.counterOverflow + } + let (nextProcessedBytes, processedBytesOverflowed) = + processedBytes.addingReportingOverflow(lineBytes) + guard !processedBytesOverflowed else { + throw RolloutTailSourceError.counterOverflow + } + processedBytes = nextProcessedBytes + processedLineCount += 1 + line.removeAll(keepingCapacity: true) + skipsCurrentLine = false + countedCurrentOversizedLine = false + completeByteOffset = lineEnd + lineStart = lineEnd + pieceStart = chunk.index(after: newline) + if !shouldContinue + || processedLineCount >= lineLimit + || processedBytes >= byteLimit { + return BoundedJSONLReadResult( + completeByteOffset: completeByteOffset, + resumeByteOffset: completeByteOffset, + bytesRead: bytesRead, + oversizedRecordCount: oversizedRecordCount, + processedLineCount: processedLineCount, + stoppedEarly: true, + discardingOversizedRecord: false + ) + } + } + let remainder = chunk[pieceStart...] + if !skipsCurrentLine { + if line.count + remainder.count > maximumRecordBytes { + line.removeAll(keepingCapacity: false) + skipsCurrentLine = true + if !countedCurrentOversizedLine { + oversizedRecordCount += 1 + countedCurrentOversizedLine = true + } + } else { + line.append(contentsOf: remainder) + } + } + if bytesRead >= byteLimit, skipsCurrentLine || !line.isEmpty { + let discardsPartial = discardsPartialRecordAtByteLimit + && !line.isEmpty + if discardsPartial, + !skipsCurrentLine, + !countedCurrentOversizedLine { + oversizedRecordCount += 1 + } + let (resumeByteOffset, offsetOverflowed) = + startOffset.addingReportingOverflow(bytesRead) + guard !offsetOverflowed else { + throw RolloutTailSourceError.counterOverflow + } + return BoundedJSONLReadResult( + completeByteOffset: completeByteOffset, + resumeByteOffset: discardsPartial || skipsCurrentLine + ? resumeByteOffset + : completeByteOffset, + bytesRead: bytesRead, + oversizedRecordCount: oversizedRecordCount, + processedLineCount: processedLineCount, + stoppedEarly: true, + discardingOversizedRecord: skipsCurrentLine + || discardsPartial + ) + } + } + let resumeByteOffset: UInt64 + if skipsCurrentLine { + let (offset, overflowed) = + startOffset.addingReportingOverflow(bytesRead) + guard !overflowed else { + throw RolloutTailSourceError.counterOverflow + } + resumeByteOffset = offset + } else { + resumeByteOffset = completeByteOffset + } + return BoundedJSONLReadResult( + completeByteOffset: completeByteOffset, + resumeByteOffset: resumeByteOffset, + bytesRead: bytesRead, + oversizedRecordCount: oversizedRecordCount, + processedLineCount: processedLineCount, + stoppedEarly: false, + discardingOversizedRecord: skipsCurrentLine + ) + } } struct RolloutFileIdentity: Codable, Equatable, Hashable, Sendable { @@ -14,6 +232,7 @@ struct RolloutCheckpoint: Codable, Equatable, Sendable { let byteLength: UInt64 let threadID: String? let fingerprint: UInt64 + let rawSuffixFingerprint: UInt64? } struct RolloutCursor: Codable, Equatable, Sendable { @@ -26,6 +245,7 @@ struct RolloutCursor: Codable, Equatable, Sendable { let threadID: String? let processedPrefixFingerprint: UInt64? let checkpoint: RolloutCheckpoint? + let discardingOversizedRecord: Bool? } enum LocalTokenComponent: String, Codable, CaseIterable, Sendable { @@ -44,9 +264,10 @@ struct LocalTokenUsage: Codable, Equatable, Sendable { let outputTokens: Int64 let reasoningOutputTokens: Int64 let totalTokens: Int64 - var observedComponents: Set = Set( - LocalTokenComponent.allCases - ) + private let observedComponentMask: UInt8 + var observedComponents: Set { + Set(LocalTokenComponent.allCases.filter(observes)) + } private enum CodingKeys: String, CodingKey { case inputTokens @@ -75,7 +296,25 @@ struct LocalTokenUsage: Codable, Equatable, Sendable { self.outputTokens = outputTokens self.reasoningOutputTokens = reasoningOutputTokens self.totalTokens = totalTokens - self.observedComponents = observedComponents + observedComponentMask = Self.mask(for: observedComponents) + } + + init( + inputTokens: Int64, + cachedInputTokens: Int64, + cacheWriteInputTokens: Int64, + outputTokens: Int64, + reasoningOutputTokens: Int64, + totalTokens: Int64, + observedComponentMask: UInt8 + ) { + self.inputTokens = inputTokens + self.cachedInputTokens = cachedInputTokens + self.cacheWriteInputTokens = cacheWriteInputTokens + self.outputTokens = outputTokens + self.reasoningOutputTokens = reasoningOutputTokens + self.totalTokens = totalTokens + self.observedComponentMask = observedComponentMask } init(from decoder: Decoder) throws { @@ -95,10 +334,12 @@ struct LocalTokenUsage: Codable, Equatable, Sendable { forKey: .reasoningOutputTokens ) totalTokens = try values.decode(Int64.self, forKey: .totalTokens) - observedComponents = try values.decodeIfPresent( + observedComponentMask = Self.mask( + for: try values.decodeIfPresent( [LocalTokenComponent].self, forKey: .observedComponents - ).map(Set.init) ?? [.total] + ).map(Set.init) ?? [.total] + ) } func encode(to encoder: Encoder) throws { @@ -120,6 +361,31 @@ struct LocalTokenUsage: Codable, Equatable, Sendable { forKey: .observedComponents ) } + + func observes(_ component: LocalTokenComponent) -> Bool { + observedComponentMask & Self.bit(for: component) != 0 + } + + func sharedObservedComponentMask(with other: LocalTokenUsage) -> UInt8 { + observedComponentMask & other.observedComponentMask + } + + private static func mask( + for components: Set + ) -> UInt8 { + components.reduce(0) { $0 | bit(for: $1) } + } + + private static func bit(for component: LocalTokenComponent) -> UInt8 { + switch component { + case .input: 1 << 0 + case .cachedInput: 1 << 1 + case .cacheWriteInput: 1 << 2 + case .output: 1 << 3 + case .reasoningOutput: 1 << 4 + case .total: 1 << 5 + } + } } struct RolloutRecord: Equatable, Sendable { @@ -156,17 +422,30 @@ struct RolloutTailBatch: Equatable, Sendable { let unsupportedRecordCount: Int let malformedRecordCount: Int let requiresRebuild: Bool + let continuityChanged: Bool + let hasMoreRecords: Bool + let processedLineCount: Int } struct IncrementalRolloutTailSource { private static let fingerprintOffsetBasis: UInt64 = 14_695_981_039_346_656_037 private static let fingerprintPrime: UInt64 = 1_099_511_628_211 + private static let payloadMarker = Data(#","payload":"#.utf8) + private static let compactedTypeMarker = Data(#""type":"compacted""#.utf8) + private static let emptyPayloadSuffix = Data(#","payload":{}}"#.utf8) + // ponytail: bound one batch; stream normalization if one-pass import matters. + static let maximumLinesPerBatch = 100_000 + static let maximumBytesPerBatch: UInt64 = 64 * 1_024 * 1_024 func read( fileURL: URL, cursor: RolloutCursor?, - observedAt: Date + observedAt: Date, + maximumLines: Int = Self.maximumLinesPerBatch, + maximumBytes: UInt64 = Self.maximumBytesPerBatch, + maximumRecordBytes: Int = BoundedJSONLReader.maximumRecordBytes ) throws -> RolloutTailBatch { + let byteBudget = max(maximumBytes, 1) let attributes = try FileManager.default.attributesOfItem(atPath: fileURL.path) guard let systemNumber = (attributes[.systemNumber] as? NSNumber)?.uint64Value, @@ -188,62 +467,68 @@ struct IncrementalRolloutTailSource { && fileSize == $0.fileSize && modificationTime == $0.modificationTime } ?? false - let continuationVerification: (matches: Bool, bytesRead: Int) + let continuationVerification: (matches: Bool?, bytesRead: UInt64) if metadataUnchanged { continuationVerification = (true, 0) - } else if sameFileIdentity { - continuationVerification = try verifiedCheckpoint( - handle: handle, - fileSize: fileSize, - cursor: cursor - ) + } else if sameFileIdentity, + let cursor, + fileSize > cursor.fileSize { + // Codex owns rollout files and appends to them. Verify the last + // complete record so live growth stays proportional to the delta. + continuationVerification = + cursor.discardingOversizedRecord == true + ? (true, 0) + : try verifiedCheckpoint( + handle: handle, + fileSize: fileSize, + cursor: cursor, + maximumBytes: byteBudget + ) } else { continuationVerification = try verifiedPrefix( handle: handle, fileSize: fileSize, - cursor: cursor + cursor: cursor, + maximumBytes: byteBudget + ) + } + if continuationVerification.matches == nil, + sameFileIdentity, + cursor.map({ fileSize > $0.fileSize }) == true, + let cursor, + cursor.checkpoint != nil { + return RolloutTailBatch( + records: [], + cursor: cursor, + bytesRead: 0, + unsupportedRecordCount: 0, + malformedRecordCount: 0, + requiresRebuild: false, + continuityChanged: false, + hasMoreRecords: true, + processedLineCount: 0 ) } - let continuesSource = continuationVerification.matches + let continuesSource = continuationVerification.matches == true let startOffset = continuesSource ? cursor?.byteOffset ?? 0 : 0 let requiresRebuild = cursor != nil && !continuesSource + let continuityChanged = cursor != nil + && continuationVerification.matches == false let sourceGeneration: UInt64 if let cursor { - sourceGeneration = sameFileIdentity && continuesSource - ? cursor.sourceGeneration - : cursor.sourceGeneration + 1 + if sameFileIdentity && continuesSource { + sourceGeneration = cursor.sourceGeneration + } else { + let (next, overflowed) = + cursor.sourceGeneration.addingReportingOverflow(1) + guard !overflowed else { + throw RolloutTailSourceError.counterOverflow + } + sourceGeneration = next + } } else { sourceGeneration = 0 } - - try handle.seek(toOffset: startOffset) - let data = try handle.readToEnd() ?? Data() - let bytesRead = UInt64(continuationVerification.bytesRead + data.count) - - guard let lastNewline = data.lastIndex(of: 0x0A) else { - return RolloutTailBatch( - records: [], - cursor: RolloutCursor( - fileIdentity: identity, - sourceGeneration: sourceGeneration, - byteOffset: startOffset, - fileSize: fileSize, - modificationTime: modificationTime, - lastOrdinal: continuesSource ? cursor?.lastOrdinal : nil, - threadID: continuesSource ? cursor?.threadID : nil, - processedPrefixFingerprint: continuesSource - ? cursor?.processedPrefixFingerprint - : nil, - checkpoint: continuesSource ? cursor?.checkpoint : nil - ), - bytesRead: bytesRead, - unsupportedRecordCount: 0, - malformedRecordCount: 0, - requiresRebuild: requiresRebuild - ) - } - - let completeData = data[...lastNewline] var threadID = continuesSource ? cursor?.threadID : nil var seenEventKeys = Set() var records: [RolloutRecord] = [] @@ -258,21 +543,30 @@ struct IncrementalRolloutTailSource { byteOffset: UInt64, byteLength: UInt64, threadID: String?, - identity: String + fingerprint: UInt64, + rawSuffixFingerprint: UInt64? )? let decoder = JSONDecoder() decoder.keyDecodingStrategy = .convertFromSnakeCase + let readByteBudget = continuationVerification.bytesRead < byteBudget + ? byteBudget - continuationVerification.bytesRead + : 1 - for line in completeData.split(separator: 0x0A, omittingEmptySubsequences: true) { - let relativeLineOffset = completeData.distance( - from: completeData.startIndex, - to: line.startIndex - ) - let absoluteLineOffset = startOffset + UInt64(relativeLineOffset) - guard let wire = try? decoder.decode(RolloutWire.self, from: Data(line)) else { + let readResult = try BoundedJSONLReader.read( + handle: handle, + from: startOffset, + maximumLines: maximumLines, + maximumBytes: readByteBudget, + maximumRecordBytes: maximumRecordBytes, + startsByDiscardingOversizedRecord: + continuesSource + && cursor?.discardingOversizedRecord == true, + discardsPartialRecordAtByteLimit: false + ) { line, absoluteLineOffset in + guard let wire = decodeWire(line, decoder: decoder) else { unsupportedRecordCount += 1 malformedRecordCount += 1 - continue + return true } let type = wire.type let payload = wire.payload @@ -299,12 +593,19 @@ struct IncrementalRolloutTailSource { let contextTokenUsage = localTokenUsage( payload.info?.lastTokenUsage ) + if totalUsage?.totalTokens != nil, tokenUsage == nil { + malformedRecordCount += 1 + } + if payload.info?.lastTokenUsage?.totalTokens != nil, + contextTokenUsage == nil { + malformedRecordCount += 1 + } let eventKey: String if let ordinal { eventKey = "\(threadID ?? "unknown")|ordinal|\(ordinal)" let isReplay = continuesSource && cursor?.lastOrdinal.map { ordinal <= $0 } == true - if isReplay { continue } + if isReplay { return true } lastObservedOrdinal = max(lastObservedOrdinal ?? 0, ordinal) } else { eventKey = [ @@ -316,25 +617,33 @@ struct IncrementalRolloutTailSource { turnID ?? "none" ].joined(separator: "|") } - guard seenEventKeys.insert(eventKey).inserted else { continue } + guard seenEventKeys.insert(eventKey).inserted else { return true } guard isSupported( type: type, eventType: eventType, toolClass: payload.item?.type ) else { unsupportedRecordCount += 1 - continue + return true } if line.count <= 4_096 { checkpointCandidate = ( byteOffset: absoluteLineOffset, byteLength: UInt64(line.count), threadID: threadID, - identity: nonContentIdentity + fingerprint: Self.fingerprint( + nonContentIdentity.utf8 + ), + rawSuffixFingerprint: nil ) } else { - checkpointCandidate = nil - checkpoint = nil + checkpointCandidate = ( + byteOffset: absoluteLineOffset, + byteLength: UInt64(line.count), + threadID: threadID, + fingerprint: Self.fingerprint(line.prefix(2_048)), + rawSuffixFingerprint: Self.fingerprint(line.suffix(2_048)) + ) } records.append( RolloutRecord( @@ -365,14 +674,29 @@ struct IncrementalRolloutTailSource { : nil ) ) + return true + } + unsupportedRecordCount += readResult.oversizedRecordCount + malformedRecordCount += readResult.oversizedRecordCount + let (bytesRead, bytesReadOverflowed) = + continuationVerification.bytesRead.addingReportingOverflow( + readResult.bytesRead + ) + guard !bytesReadOverflowed else { + throw RolloutTailSourceError.counterOverflow + } + let byteOffset = readResult.resumeByteOffset + let hasMoreRecords = readResult.stoppedEarly && byteOffset < fileSize + if byteOffset == startOffset, !continuesSource { + processedPrefixFingerprint = nil } - let byteOffset = startOffset + UInt64(completeData.count) if let candidate = checkpointCandidate { checkpoint = RolloutCheckpoint( byteOffset: candidate.byteOffset, byteLength: candidate.byteLength, threadID: candidate.threadID, - fingerprint: Self.fingerprint(candidate.identity.utf8) + fingerprint: candidate.fingerprint, + rawSuffixFingerprint: candidate.rawSuffixFingerprint ) } @@ -387,62 +711,135 @@ struct IncrementalRolloutTailSource { lastOrdinal: lastObservedOrdinal, threadID: threadID, processedPrefixFingerprint: processedPrefixFingerprint, - checkpoint: checkpoint + checkpoint: checkpoint, + discardingOversizedRecord: + readResult.discardingOversizedRecord ? true : nil ), bytesRead: bytesRead, unsupportedRecordCount: unsupportedRecordCount, malformedRecordCount: malformedRecordCount, - requiresRebuild: requiresRebuild + requiresRebuild: requiresRebuild, + continuityChanged: continuityChanged, + hasMoreRecords: hasMoreRecords, + processedLineCount: readResult.processedLineCount ) } private func verifiedPrefix( handle: FileHandle, fileSize: UInt64, - cursor: RolloutCursor? - ) throws -> (matches: Bool, bytesRead: Int) { + cursor: RolloutCursor?, + maximumBytes: UInt64 + ) throws -> (matches: Bool?, bytesRead: UInt64) { guard let cursor, cursor.byteOffset <= fileSize, - cursor.byteOffset <= UInt64(Int.max), let expectedFingerprint = cursor.processedPrefixFingerprint else { return (false, 0) } - try handle.seek(toOffset: 0) - let prefix = try handle.read(upToCount: Int(cursor.byteOffset)) ?? Data() + guard cursor.byteOffset <= maximumBytes else { + return (nil, 0) + } + var fingerprint = Self.fingerprintOffsetBasis + var threadID: String? + let decoder = JSONDecoder() + decoder.keyDecodingStrategy = .convertFromSnakeCase + let result = try BoundedJSONLReader.read( + handle: handle, + from: 0, + through: cursor.byteOffset + ) { line, absoluteLineOffset in + guard let wire = decodeWire(line, decoder: decoder) else { + return true + } + if wire.type == "session_meta", + let observedThreadID = wire.payload.id { + threadID = observedThreadID + } + fingerprint = Self.fingerprint( + replayIdentity( + wire: wire, + threadID: threadID, + absoluteLineOffset: absoluteLineOffset + ).utf8, + seed: fingerprint + ) + return true + } return ( - prefix.count == Int(cursor.byteOffset) - && replayFingerprint(prefix) == expectedFingerprint, - prefix.count + result.completeByteOffset == cursor.byteOffset + && fingerprint == expectedFingerprint, + result.bytesRead ) } private func verifiedCheckpoint( handle: FileHandle, fileSize: UInt64, - cursor: RolloutCursor? - ) throws -> (matches: Bool, bytesRead: Int) { + cursor: RolloutCursor?, + maximumBytes: UInt64 + ) throws -> (matches: Bool?, bytesRead: UInt64) { guard let cursor else { return (true, 0) } guard cursor.byteOffset > 0 else { return (true, 0) } + let checkpointEnd = cursor.checkpoint.map { + $0.byteOffset.addingReportingOverflow($0.byteLength) + } guard let checkpoint = cursor.checkpoint, - checkpoint.byteLength <= 4_096, - checkpoint.byteOffset + checkpoint.byteLength <= fileSize + let checkpointEnd, + !checkpointEnd.overflow, + checkpointEnd.partialValue <= fileSize else { return (false, 0) } + if let expectedSuffix = checkpoint.rawSuffixFingerprint { + guard checkpoint.byteLength <= UInt64( + BoundedJSONLReader.maximumRecordBytes + ) else { + return (false, 0) + } + let sliceLength = min(checkpoint.byteLength, 2_048) + guard sliceLength * 2 <= maximumBytes else { + return (nil, 0) + } + try handle.seek(toOffset: checkpoint.byteOffset) + let prefix = try handle.read( + upToCount: Int(sliceLength) + ) ?? Data() + let suffixOffset = checkpointEnd.partialValue - sliceLength + try handle.seek(toOffset: suffixOffset) + let suffix = try handle.read( + upToCount: Int(sliceLength) + ) ?? Data() + let bytesRead = UInt64(prefix.count + suffix.count) + guard prefix.count == Int(sliceLength), + suffix.count == Int(sliceLength) else { + return (false, bytesRead) + } + return ( + Self.fingerprint(prefix) == checkpoint.fingerprint + && Self.fingerprint(suffix) == expectedSuffix, + bytesRead + ) + } + guard checkpoint.byteLength <= 4_096 else { + return (false, 0) + } + guard checkpoint.byteLength <= maximumBytes else { + return (nil, 0) + } try handle.seek(toOffset: checkpoint.byteOffset) let data = try handle.read( upToCount: Int(checkpoint.byteLength) ) ?? Data() guard data.count == Int(checkpoint.byteLength) else { - return (false, data.count) + return (false, UInt64(data.count)) } let decoder = JSONDecoder() decoder.keyDecodingStrategy = .convertFromSnakeCase - guard let wire = try? decoder.decode(RolloutWire.self, from: data) else { - return (false, data.count) + guard let wire = decodeWire(data, decoder: decoder) else { + return (false, UInt64(data.count)) } let threadID = wire.type == "session_meta" ? wire.payload.id @@ -454,37 +851,10 @@ struct IncrementalRolloutTailSource { ) return ( Self.fingerprint(identity.utf8) == checkpoint.fingerprint, - data.count + UInt64(data.count) ) } - private func replayFingerprint(_ data: Data) -> UInt64 { - var fingerprint = Self.fingerprintOffsetBasis - var threadID: String? - let decoder = JSONDecoder() - decoder.keyDecodingStrategy = .convertFromSnakeCase - for line in data.split(separator: 0x0A, omittingEmptySubsequences: true) { - guard let wire = try? decoder.decode(RolloutWire.self, from: Data(line)) else { - continue - } - if wire.type == "session_meta", let observedThreadID = wire.payload.id { - threadID = observedThreadID - } - let absoluteLineOffset = UInt64( - data.distance(from: data.startIndex, to: line.startIndex) - ) - fingerprint = Self.fingerprint( - replayIdentity( - wire: wire, - threadID: threadID, - absoluteLineOffset: absoluteLineOffset - ).utf8, - seed: fingerprint - ) - } - return fingerprint - } - private func replayIdentity( wire: RolloutWire, threadID: String?, @@ -542,10 +912,38 @@ struct IncrementalRolloutTailSource { return components.joined(separator: "|") } + private func decodeWire( + _ data: Data, + decoder: JSONDecoder + ) -> RolloutWire? { + if let payload = data.range(of: Self.payloadMarker), + data.range( + of: Self.compactedTypeMarker, + in: data.startIndex.. LocalTokenUsage? { - usage?.totalTokens.map { + usage?.totalTokens.flatMap { + let counters = [ + usage?.inputTokens, + usage?.cachedInputTokens, + usage?.cacheWriteInputTokens, + usage?.outputTokens, + usage?.reasoningOutputTokens, + usage?.totalTokens + ].compactMap { $0 } + guard counters.allSatisfy({ $0 >= 0 }) else { return nil } var observed: Set = [.total] if usage?.inputTokens != nil { observed.insert(.input) } if usage?.cachedInputTokens != nil { diff --git a/Sources/CodexLimits/SettingsView.swift b/Sources/CodexLimits/SettingsView.swift index 04b2658..e4f05ee 100644 --- a/Sources/CodexLimits/SettingsView.swift +++ b/Sources/CodexLimits/SettingsView.swift @@ -11,8 +11,17 @@ struct SettingsView: View { var body: some View { Form { - Stepper(value: $safetyBuffer, in: 1 ... 10, step: 1) { - Text("Safety buffer: \(Int(safetyBuffer))%") + Stepper( + value: Binding( + get: { SafetyBufferPolicy.normalized(safetyBuffer) }, + set: { safetyBuffer = SafetyBufferPolicy.normalized($0) } + ), + in: SafetyBufferPolicy.range, + step: 1 + ) { + Text( + "Safety buffer: \(Int(SafetyBufferPolicy.normalized(safetyBuffer)))%" + ) } .onChange(of: safetyBuffer) { _, value in monitor.updateSafetyBuffer(value) diff --git a/Sources/CodexLimits/UsageHistory.swift b/Sources/CodexLimits/UsageHistory.swift index 1dcb0f0..d0b31b4 100644 --- a/Sources/CodexLimits/UsageHistory.swift +++ b/Sources/CodexLimits/UsageHistory.swift @@ -142,6 +142,7 @@ actor UsageHistory { private static let accountBindingName = ".codex-limits-account.json" private static let maximumFileSize = 1_000_000 private static let maximumGeneration = 1_000_000_000 + private static let automaticSyncInterval: TimeInterval = 30 * 60 private let localDirectory: URL private let installationID: String @@ -152,6 +153,7 @@ actor UsageHistory { private var deletionStatus: DeletionStatus = .none private var migrationWarning = false private var syncAccountBindingToken: String? + private var lastSynchronizationAttemptAt: Date? private let beforeCoordinatedMarkerRead: ((URL) throws -> Void)? init( @@ -182,7 +184,10 @@ actor UsageHistory { } catch { errorMessage = "Usage history couldn’t be saved." } - return state(fallback: legacySamples) + return state( + fallback: legacySamples, + refreshSamples: true + ) } func restoreExistingState() -> State? { @@ -236,6 +241,7 @@ actor UsageHistory { installationID: installationID, coordinated: false ) + knownSamples = normalized(knownSamples + [sample]) if let syncDirectory { try prepareRoot( syncDirectory, @@ -304,7 +310,7 @@ actor UsageHistory { } syncDirectory = directory errorMessage = nil - return state() + return state(refreshSamples: true) } if let accountIdentity { syncAccountBindingToken = try ensureAccountBinding( @@ -348,6 +354,7 @@ actor UsageHistory { syncDirectory = nil syncAccountBindingToken = nil errorMessage = nil + lastSynchronizationAttemptAt = nil return state() } @@ -364,13 +371,29 @@ actor UsageHistory { self.partition = partition syncDirectory = nil syncAccountBindingToken = nil + lastSynchronizationAttemptAt = nil errorMessage = nil knownSamples = [] return load() } func synchronize() -> State { - guard let syncDirectory else { return state() } + synchronize(at: Date()) + } + + func synchronizeIfDue(at now: Date = Date()) -> State { + if let lastSynchronizationAttemptAt { + let elapsed = now.timeIntervalSince(lastSynchronizationAttemptAt) + if elapsed >= 0, elapsed < Self.automaticSyncInterval { + return state() + } + } + return synchronize(at: now) + } + + private func synchronize(at now: Date) -> State { + lastSynchronizationAttemptAt = now + guard let syncDirectory else { return state(refreshSamples: true) } do { try prepareLocalStore() try prepareRoot( @@ -389,7 +412,7 @@ actor UsageHistory { generation: localMarker.generation ?? 1 ) errorMessage = nil - return state() + return state(refreshSamples: true) } try reconcileGeneration(with: syncDirectory) let generation = try effectiveGeneration(in: syncDirectory) @@ -404,7 +427,7 @@ actor UsageHistory { } catch { errorMessage = message(for: error) } - return state() + return state(refreshSamples: true) } func deleteAnalyticsHistory( @@ -496,7 +519,7 @@ actor UsageHistory { : "Deletion pending — sync folder unavailable." deletionStatus = deletionTarget == nil ? .none : .pendingSync } - return state() + return state(refreshSamples: true) } func retryPendingDeletion( @@ -546,7 +569,7 @@ actor UsageHistory { } deletionStatus = .pendingSync } - return state() + return state(refreshSamples: true) } func rebuildAvailableHistory(_ samples: [UsageSample]) -> State { @@ -589,23 +612,27 @@ actor UsageHistory { } catch { errorMessage = "Available history couldn’t be rebuilt." } - return state() + return state(refreshSamples: true) } private func state( fallback: [UsageSample] = [], - issue: Issue? = nil + issue: Issue? = nil, + refreshSamples: Bool = false ) -> State { - let local = readAll(from: activeLocalDirectory) - if local.hadError && errorMessage == nil { - errorMessage = "Some usage history couldn’t be read." + if refreshSamples { + let local = readAll(from: activeLocalDirectory) + if local.hadError && errorMessage == nil { + errorMessage = "Some usage history couldn’t be read." + } + knownSamples = local.hadError || errorMessage != nil + ? normalized(local.samples + knownSamples + fallback) + : local.samples + } else if !fallback.isEmpty { + knownSamples = normalized(knownSamples + fallback) } - let samples = local.hadError || errorMessage != nil - ? normalized(local.samples + knownSamples + fallback) - : local.samples - knownSamples = samples return State( - samples: samples, + samples: knownSamples, folderName: syncDirectory?.lastPathComponent, errorMessage: errorMessage, deletionStatus: deletionStatus, @@ -937,7 +964,11 @@ actor UsageHistory { guard (size ?? 0) <= Self.maximumFileSize else { throw HistoryError.invalidFile } - return try Data(contentsOf: url) + let data = try Data(contentsOf: url) + guard data.count <= Self.maximumFileSize else { + throw HistoryError.invalidFile + } + return data } private func isUbiquitousItem(_ url: URL) -> Bool { @@ -952,12 +983,7 @@ actor UsageHistory { } private func normalized(_ samples: [UsageSample]) -> [UsageSample] { - let valid = samples.filter { - $0.observedAt <= $0.resetsAt - && $0.remainingPercent.isFinite - && (0 ... 100).contains($0.remainingPercent) - && ($0.lifetimeTokens.map { $0 >= 0 } ?? true) - } + let valid = samples.filter(\.isValid) var byIdentity: [UsageSample: UsageSample] = [:] for sample in valid { guard let existing = byIdentity[sample] else { diff --git a/Sources/CodexLimits/UsageIntelligenceEngine.swift b/Sources/CodexLimits/UsageIntelligenceEngine.swift index ab85fd1..3535d5b 100644 --- a/Sources/CodexLimits/UsageIntelligenceEngine.swift +++ b/Sources/CodexLimits/UsageIntelligenceEngine.swift @@ -397,6 +397,8 @@ struct UsageIntelligenceInput: Equatable, Sendable { let localActivityHistoryFacts: [LocalActivityFact] let localActivityObservation: LocalActivityObservation let localTaskProjections: [ThreadProjection] + let localActivityContentRevision: UInt64? + let reusableLocalAggregates: LocalAggregateCache? let compatibleTokenSources: Set let analyticsExploration: AnalyticsExplorationState let insightDispositions: [String: InsightDisposition] @@ -416,6 +418,8 @@ struct UsageIntelligenceInput: Equatable, Sendable { "Codex local records are unavailable" ), localTaskProjections: [ThreadProjection] = [], + localActivityContentRevision: UInt64? = nil, + reusableLocalAggregates: LocalAggregateCache? = nil, compatibleTokenSources: Set = [], analyticsExploration: AnalyticsExplorationState = .initial, insightDispositions: [String: InsightDisposition] = [:] @@ -433,12 +437,115 @@ struct UsageIntelligenceInput: Equatable, Sendable { localActivityHistoryFacts ?? localActivityFacts self.localActivityObservation = localActivityObservation self.localTaskProjections = localTaskProjections + self.localActivityContentRevision = localActivityContentRevision + self.reusableLocalAggregates = reusableLocalAggregates self.compatibleTokenSources = compatibleTokenSources self.analyticsExploration = analyticsExploration self.insightDispositions = insightDispositions } } +struct LocalAggregateCache: Equatable, Sendable { + let contentRevision: UInt64 + let observation: LocalActivityObservation + let workloadMixChanged: Bool + let localTokenActivity: LocalTokenActivitySnapshot + let usageReceipts: UsageReceiptSnapshot + let activityTimeline: ActivityTimelineSnapshot + fileprivate let localHistory: LocalHistoryAggregateCache? + + static func == (lhs: Self, rhs: Self) -> Bool { + lhs.contentRevision == rhs.contentRevision + && lhs.observation == rhs.observation + && lhs.workloadMixChanged == rhs.workloadMixChanged + && lhs.localTokenActivity == rhs.localTokenActivity + && lhs.usageReceipts == rhs.usageReceipts + && lhs.activityTimeline == rhs.activityTimeline + } +} + +private final class LocalHistoryAggregateCache: @unchecked Sendable { + let factIndex: LocalActivityFactIndex + let samples: [UsageSample] + let observation: LocalActivityObservation + let accountPartitionID: String + let limitID: String + let currentReset: Date + let compatibleTokenSources: Set + let weeklyEvidence: WeeklyUsageEvidenceSet + let activeTimeHistory: ActiveTimeHistorySelection + + init( + factIndex: LocalActivityFactIndex, + samples: [UsageSample], + observation: LocalActivityObservation, + accountPartitionID: String, + limitID: String, + currentReset: Date, + compatibleTokenSources: Set, + weeklyEvidence: WeeklyUsageEvidenceSet, + activeTimeHistory: ActiveTimeHistorySelection + ) { + self.factIndex = factIndex + self.samples = samples + self.observation = observation + self.accountPartitionID = accountPartitionID + self.limitID = limitID + self.currentReset = currentReset + self.compatibleTokenSources = compatibleTokenSources + self.weeklyEvidence = weeklyEvidence + self.activeTimeHistory = activeTimeHistory + } + + func matches( + samples: [UsageSample], + observation: LocalActivityObservation, + accountPartitionID: String, + limitID: String, + currentReset: Date, + compatibleTokenSources: Set + ) -> Bool { + self.accountPartitionID == accountPartitionID + && self.limitID == limitID + && self.currentReset == currentReset + && self.compatibleTokenSources == compatibleTokenSources + && observationsHaveSameHistoryEffect( + self.observation, + observation + ) + && samplesMatchExactly(self.samples, samples) + } + + private func observationsHaveSameHistoryEffect( + _ lhs: LocalActivityObservation, + _ rhs: LocalActivityObservation + ) -> Bool { + switch (lhs, rhs) { + case let ( + .continuous(lhsVersion, _), + .continuous(rhsVersion, _) + ): + lhsVersion == rhsVersion + default: + lhs == rhs + } + } + + private func samplesMatchExactly( + _ lhs: [UsageSample], + _ rhs: [UsageSample] + ) -> Bool { + lhs.count == rhs.count + && zip(lhs, rhs).allSatisfy { + $0.observedAt == $1.observedAt + && $0.remainingPercent == $1.remainingPercent + && $0.resetsAt == $1.resetsAt + && $0.lifetimeTokens == $1.lifetimeTokens + && $0.comparisonBreak == $1.comparisonBreak + } + } +} + struct UsageReaderSnapshot: Equatable, Sendable { let account: UsageSnapshot? let accountSource: UsageValueSource @@ -458,6 +565,7 @@ struct UsageReaderSnapshot: Equatable, Sendable { let activeTimeAvailability: ActiveTimeAvailabilitySnapshot let localTaskProjections: [ThreadProjection] let accountPartitionID: String? + let reusableLocalAggregates: LocalAggregateCache? var insights: DeterministicInsightsSnapshot var fetchedAt: Date? { account?.fetchedAt } @@ -491,13 +599,18 @@ struct UsageReaderSnapshot: Equatable, Sendable { func updatedText(at now: Date) -> String { guard let fetchedAt = account?.fetchedAt else { return "Not updated" } - let seconds = max(now.timeIntervalSince(fetchedAt), 0) + let rawSeconds = now.timeIntervalSince(fetchedAt) + guard rawSeconds.isFinite else { return "Updated a long time ago" } + let seconds = max(rawSeconds, 0) if seconds < 60 { return "Updated just now" } if seconds < 3_600 { return "Updated \(Int(seconds / 60)) min ago" } if seconds < 86_400 { let hours = Int(seconds / 3_600) return "Updated \(hours) \(hours == 1 ? "hr" : "hrs") ago" } + guard seconds / 86_400 <= Double(Int.max) else { + return "Updated a long time ago" + } let days = Int(seconds / 86_400) return "Updated \(days) \(days == 1 ? "day" : "days") ago" } @@ -534,12 +647,26 @@ enum UsageIntelligenceEngine { .sorted { $0.observedAt < $1.observedAt } } } ?? [] + let reusableLocalAggregates = input.reusableLocalAggregates.flatMap { + cache -> LocalAggregateCache? in + guard let revision = input.localActivityContentRevision, + revision != 0, + cache.contentRevision == revision, + canReuseLocalAggregates( + from: cache.observation, + to: input.localActivityObservation + ) else { + return nil + } + return cache + } let workloadMixChanged = input.account?.mainLimit.map { - LocalWorkloadMixAnalyzer.detectsChange( - facts: input.localActivityFacts, - observation: input.localActivityObservation, - window: $0.window - ) + reusableLocalAggregates?.workloadMixChanged + ?? LocalWorkloadMixAnalyzer.detectsChange( + facts: input.localActivityFacts, + observation: input.localActivityObservation, + window: $0.window + ) } ?? false let evidence = evidence( account: input.account, @@ -617,47 +744,109 @@ enum UsageIntelligenceEngine { accountActivity: accountTokenActivity, accountEpochStartedAt: input.accountEpochStartedAt ) { - localTokenActivity = LocalTokenActivityAggregator.evaluate( - facts: input.localActivityFacts, - interval: interval, - observation: input.localActivityObservation - ) + if let cached = reusableLocalAggregates?.localTokenActivity, + cached.interval.start == interval.start, + !tokenFactsAffectIntervalChange( + input.localActivityFacts, + from: cached.interval, + to: interval + ) { + localTokenActivity = cached.updating( + interval: interval, + observation: input.localActivityObservation + ) + } else { + localTokenActivity = LocalTokenActivityAggregator.evaluate( + facts: input.localActivityFacts, + interval: interval, + observation: input.localActivityObservation + ) + } } else { localTokenActivity = .unavailable( "Weekly token interval is unavailable", interval: DateInterval(start: input.now, end: input.now) ) } - let usageReceipts = UsageReceiptAggregator.evaluate( - facts: input.localActivityFacts, - projections: input.localTaskProjections, - interval: localTokenActivity.interval, - observation: input.localActivityObservation - ) - let activityTimeline = ActivityTimelineAggregator.evaluate( - facts: input.localActivityFacts, - projections: input.localTaskProjections, - interval: localTokenActivity.interval, - observation: input.localActivityObservation - ) + let usageReceipts = reusableLocalAggregates.flatMap { + $0.usageReceipts.interval.start + == localTokenActivity.interval.start + ? $0.usageReceipts.updating( + interval: localTokenActivity.interval, + observation: input.localActivityObservation + ) + : nil + } + ?? UsageReceiptAggregator.evaluate( + facts: input.localActivityFacts, + projections: input.localTaskProjections, + interval: localTokenActivity.interval, + observation: input.localActivityObservation + ) + let activityTimeline = reusableLocalAggregates.flatMap { + $0.activityTimeline.interval.start + == localTokenActivity.interval.start + ? $0.activityTimeline.updating( + interval: localTokenActivity.interval, + observation: input.localActivityObservation + ) + : nil + } + ?? ActivityTimelineAggregator.evaluate( + facts: input.localActivityFacts, + projections: input.localTaskProjections, + interval: localTokenActivity.interval, + observation: input.localActivityObservation + ) + let reusableLocalHistory = reusableLocalAggregates?.localHistory + let localHistoryFactIndex: LocalActivityFactIndex? + if input.accountPartitionID != nil, + input.account?.mainLimit != nil { + localHistoryFactIndex = reusableLocalHistory?.factIndex + ?? LocalActivityFactIndex(input.localActivityHistoryFacts) + } else { + localHistoryFactIndex = nil + } + let weeklyEvidence: WeeklyUsageEvidenceSet? + let reusedWeeklyEvidence: Bool let sourceUsagePerToken: UsagePerTokenSnapshot if let partitionID = input.accountPartitionID, - let weekly = input.account?.mainLimit { - let evidence = WeeklyUsageEvidenceBuilder.build( - samples: input.samples, - localFacts: input.localActivityHistoryFacts, - localObservation: input.localActivityObservation, - accountPartitionID: partitionID, - limitID: weekly.limitId, - currentReset: weekly.window.resetsAt, - compatibleTokenSources: input.compatibleTokenSources - ) + let weekly = input.account?.mainLimit, + let localHistoryFactIndex { + let evidence: WeeklyUsageEvidenceSet + if let reusableLocalHistory, + reusableLocalHistory.matches( + samples: input.samples, + observation: input.localActivityObservation, + accountPartitionID: partitionID, + limitID: weekly.limitId, + currentReset: weekly.window.resetsAt, + compatibleTokenSources: input.compatibleTokenSources + ) { + evidence = reusableLocalHistory.weeklyEvidence + reusedWeeklyEvidence = true + } else { + evidence = WeeklyUsageEvidenceBuilder.build( + samples: input.samples, + localFacts: input.localActivityHistoryFacts, + localObservation: input.localActivityObservation, + accountPartitionID: partitionID, + limitID: weekly.limitId, + currentReset: weekly.window.resetsAt, + compatibleTokenSources: input.compatibleTokenSources, + factIndex: localHistoryFactIndex + ) + reusedWeeklyEvidence = false + } + weeklyEvidence = evidence sourceUsagePerToken = UsagePerTokenEngine.evaluate( current: evidence.current, history: evidence.history, pinnedBaselineID: nil ) } else { + weeklyEvidence = nil + reusedWeeklyEvidence = false sourceUsagePerToken = UsagePerTokenEngine.evaluate( current: nil, history: [], @@ -679,16 +868,27 @@ enum UsageIntelligenceEngine { in: activityTimeline.interval, filters: .all ) - let activeTimeHistory = ActiveTimeWeekEvidenceBuilder.build( - currentUsage: usagePerToken.current, - usage: usagePerToken.history, - facts: input.localActivityHistoryFacts, - projections: input.localTaskProjections, - observation: input.localActivityObservation - ) + let activeTimeHistory: ActiveTimeHistorySelection + if reusedWeeklyEvidence, let reusableLocalHistory { + activeTimeHistory = reusableLocalHistory.activeTimeHistory + } else { + activeTimeHistory = ActiveTimeWeekEvidenceBuilder.build( + currentUsage: usagePerToken.current, + usage: usagePerToken.history, + facts: input.localActivityHistoryFacts, + projections: input.localTaskProjections, + observation: input.localActivityObservation, + factIndex: localHistoryFactIndex + ) + } let activeTimeAvailability = ActiveTimeAvailabilityEngine.evaluate( currentUsage: usagePerToken.current, activeTimeThisWeek: activeTimeSlice.activeTime, + maximumConcurrency: activeTimeSlice.maximumConcurrency, + waitingTime: activeTimeSlice.waitingTime, + pollingTime: activeTimeSlice.pollingTime, + activityBreakdownReason: + activeTimeSlice.activityBreakdownReason, activeTimeCoverage: activeTimeSlice.coverage, activeTimeReason: activeTimeSlice.reason, history: activeTimeHistory.evidence, @@ -696,6 +896,40 @@ enum UsageIntelligenceEngine { usageRemainingPercent: input.account?.mainLimit?.window.remainingPercent ?? .nan ) + let localHistoryCache: LocalHistoryAggregateCache? + if reusedWeeklyEvidence, let reusableLocalHistory { + localHistoryCache = reusableLocalHistory + } else if let localHistoryFactIndex, + let weeklyEvidence, + let partitionID = input.accountPartitionID, + let weekly = input.account?.mainLimit { + localHistoryCache = LocalHistoryAggregateCache( + factIndex: localHistoryFactIndex, + samples: input.samples, + observation: input.localActivityObservation, + accountPartitionID: partitionID, + limitID: weekly.limitId, + currentReset: weekly.window.resetsAt, + compatibleTokenSources: input.compatibleTokenSources, + weeklyEvidence: weeklyEvidence, + activeTimeHistory: activeTimeHistory + ) + } else { + localHistoryCache = nil + } + let localAggregateCache: LocalAggregateCache? = + input.localActivityContentRevision.flatMap { revision in + guard revision != 0 else { return nil } + return LocalAggregateCache( + contentRevision: revision, + observation: input.localActivityObservation, + workloadMixChanged: workloadMixChanged, + localTokenActivity: localTokenActivity, + usageReceipts: usageReceipts, + activityTimeline: activityTimeline, + localHistory: localHistoryCache + ) + } let observedInterval = input.account.flatMap { account in account.mainLimit.map { UsageObservedInterval( @@ -752,10 +986,46 @@ enum UsageIntelligenceEngine { activeTimeAvailability: activeTimeAvailability, localTaskProjections: input.localTaskProjections, accountPartitionID: input.accountPartitionID, + reusableLocalAggregates: localAggregateCache, insights: insights ) } + private static func tokenFactsAffectIntervalChange( + _ facts: [LocalActivityFact], + from previous: DateInterval, + to current: DateInterval + ) -> Bool { + guard previous != current else { return false } + let parser = LocalEventTimestampParser() + return facts.contains { fact in + guard fact.key == .token, + fact.availability == .available, + let timestamp = fact.eventTimestamp, + let date = parser.date(from: timestamp) else { + return false + } + let wasIncluded = date >= previous.start && date < previous.end + let isIncluded = date >= current.start && date < current.end + return wasIncluded != isIncluded + } + } + + private static func canReuseLocalAggregates( + from previous: LocalActivityObservation, + to current: LocalActivityObservation + ) -> Bool { + switch (previous, current) { + case (.continuous, .continuous), + (.continuous, .gap), + (.gap, .gap), + (.unavailable, .unavailable): + true + default: + false + } + } + private static func bankedResetSummary( account: UsageSnapshot?, freshness: UsageFreshness, diff --git a/Sources/CodexLimits/UsageModels.swift b/Sources/CodexLimits/UsageModels.swift index a6f6313..f9c5e7d 100644 --- a/Sources/CodexLimits/UsageModels.swift +++ b/Sources/CodexLimits/UsageModels.swift @@ -1,15 +1,67 @@ import Foundation +extension Date { + var isSupportedUsageDate: Bool { + timeIntervalSinceReferenceDate.isFinite + && self >= .distantPast + && self <= .distantFuture + } +} + struct UsageWindow: Codable, Equatable, Sendable { let remainingPercent: Double let resetsAt: Date let durationMinutes: Int + var isValid: Bool { + let duration = Double(durationMinutes) * 60 + return remainingPercent.isFinite + && (0 ... 100).contains(remainingPercent) + && resetsAt.isSupportedUsageDate + && durationMinutes > 0 + && duration.isFinite + && resetsAt.addingTimeInterval(-duration).isSupportedUsageDate + } + var startsAt: Date { resetsAt.addingTimeInterval(-Double(durationMinutes) * 60) } } +extension UsageWindow { + private enum CodingKeys: String, CodingKey { + case remainingPercent + case resetsAt + case durationMinutes + } + + init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + let remainingPercent = try container.decode( + Double.self, + forKey: .remainingPercent + ) + let resetsAt = try container.decode(Date.self, forKey: .resetsAt) + let durationMinutes = try container.decode( + Int.self, + forKey: .durationMinutes + ) + self.init( + remainingPercent: remainingPercent, + resetsAt: resetsAt, + durationMinutes: durationMinutes + ) + guard isValid else { + throw DecodingError.dataCorrupted( + .init( + codingPath: decoder.codingPath, + debugDescription: "Invalid usage window." + ) + ) + } + } +} + struct UsageSample: Codable, Equatable, Hashable, Sendable { let observedAt: Date let remainingPercent: Double @@ -17,6 +69,15 @@ struct UsageSample: Codable, Equatable, Hashable, Sendable { let lifetimeTokens: Int64? let comparisonBreak: Bool + var isValid: Bool { + observedAt.isSupportedUsageDate + && resetsAt.isSupportedUsageDate + && observedAt <= resetsAt + && remainingPercent.isFinite + && (0 ... 100).contains(remainingPercent) + && (lifetimeTokens.map { $0 >= 0 } ?? true) + } + private enum CodingKeys: String, CodingKey { case observedAt case date @@ -161,6 +222,12 @@ struct AccountSpendControlFacts: Codable, Equatable, Sendable { let resetsAt: Date let reached: Bool? + var isValid: Bool { + remainingPercent.isFinite + && (0 ... 100).contains(remainingPercent) + && resetsAt.isSupportedUsageDate + } + func fillingMissingValues( from previous: AccountSpendControlFacts ) -> AccountSpendControlFacts { @@ -174,6 +241,41 @@ struct AccountSpendControlFacts: Codable, Equatable, Sendable { } } +extension AccountSpendControlFacts { + private enum CodingKeys: String, CodingKey { + case limit + case used + case remainingPercent + case resetsAt + case reached + } + + init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + self.init( + limit: try container.decode(String.self, forKey: .limit), + used: try container.decode(String.self, forKey: .used), + remainingPercent: try container.decode( + Double.self, + forKey: .remainingPercent + ), + resetsAt: try container.decode(Date.self, forKey: .resetsAt), + reached: try container.decodeIfPresent( + Bool.self, + forKey: .reached + ) + ) + guard isValid else { + throw DecodingError.dataCorrupted( + .init( + codingPath: decoder.codingPath, + debugDescription: "Invalid spend control." + ) + ) + } + } +} + struct AccountFacts: Codable, Equatable, Sendable { let lifetimeTokens: Int64? let lifetimeTokensObservedAt: Date? @@ -238,6 +340,28 @@ struct AccountFacts: Codable, Equatable, Sendable { && spendControl == nil } + var isValid: Bool { + [ + lifetimeTokens, + peakDailyTokens, + longestRunningTurnSeconds, + currentStreakDays, + longestStreakDays + ].compactMap { $0 }.allSatisfy { $0 >= 0 } + && [ + lifetimeTokensObservedAt, + peakDailyTokensObservedAt, + longestRunningTurnObservedAt, + currentStreakObservedAt, + longestStreakObservedAt, + creditsObservedAt, + creditBalanceObservedAt, + spendControlObservedAt, + spendControlReachedObservedAt + ].compactMap { $0 }.allSatisfy(\.isSupportedUsageDate) + && (spendControl?.isValid ?? true) + } + func fillingMissingValues(from previous: AccountFacts) -> AccountFacts { AccountFacts( lifetimeTokens: lifetimeTokens ?? previous.lifetimeTokens, @@ -314,6 +438,21 @@ struct UsageSnapshot: Codable, Equatable, Sendable { self.fetchedAt = fetchedAt self.accountFacts = accountFacts } + + var isValid: Bool { + fetchedAt.isSupportedUsageDate + && emergencyResetCount >= 0 + && (mainLimit?.window.isValid ?? true) + && otherLimits.allSatisfy(\.window.isValid) + && tokenHistory.allSatisfy { + $0.date.isSupportedUsageDate && $0.tokens >= 0 + } + && (bankedResetDetails ?? []).allSatisfy { + $0.expiresAt.isSupportedUsageDate + && ($0.grantedAt?.isSupportedUsageDate ?? true) + } + && (accountFacts?.isValid ?? true) + } } enum PaceStatus: String, Codable, Equatable, Sendable { diff --git a/Sources/CodexLimits/UsageMonitor.swift b/Sources/CodexLimits/UsageMonitor.swift index 632ca24..ec6458a 100644 --- a/Sources/CodexLimits/UsageMonitor.swift +++ b/Sources/CodexLimits/UsageMonitor.swift @@ -2,6 +2,16 @@ import AppKit import Combine import Foundation +enum SafetyBufferPolicy { + static let defaultValue = 3.0 + static let range = 1.0 ... 10.0 + + static func normalized(_ value: Double?) -> Double { + guard let value, value.isFinite else { return defaultValue } + return min(max(value, range.lowerBound), range.upperBound) + } +} + @MainActor final class UsageMonitor: ObservableObject { static let safetyBufferKey = "safetyBuffer" @@ -42,6 +52,8 @@ final class UsageMonitor: ObservableObject { private let history: UsageHistory private let codexAssistedHistory: CodexAssistedHistory? private let fetchUsage: () async throws -> CodexFetchResult + private let evaluateUsage: + @Sendable (UsageIntelligenceInput) -> UsageReaderSnapshot private let localActivityCollector: LocalActivityCollector? private let resetReminderCoordinator: ResetReminderCoordinator private var historyPartition: AccountHistoryPartition @@ -64,6 +76,10 @@ final class UsageMonitor: ObservableObject { private var localActivityCollection = LocalActivityCollection.unavailable( "Codex local records are unavailable" ) + private var evaluationGeneration: UInt64 = 0 + private var evaluationTask: Task? + private var localImportGeneration: UInt64 = 0 + private var localImportTask: Task? convenience init() { self.init( @@ -90,12 +106,25 @@ final class UsageMonitor: ObservableObject { resetReminderScheduler: (any ResetReminderScheduling)? = nil, resetReminderNow: @escaping () -> Date = Date.init, codexAssistedHistory: CodexAssistedHistory? = nil, - fetchUsage: @escaping () async throws -> CodexFetchResult = CodexClient.fetch + fetchUsage: @escaping () async throws -> CodexFetchResult = CodexClient.fetch, + evaluateUsage: @escaping @Sendable ( + UsageIntelligenceInput + ) -> UsageReaderSnapshot = { + UsageIntelligenceEngine.evaluate($0) + } ) { self.defaults = defaults self.fetchUsage = fetchUsage + self.evaluateUsage = evaluateUsage self.localActivityCollector = localActivityCollector self.codexAssistedHistory = codexAssistedHistory + let storedSafetyBuffer = defaults.object( + forKey: Self.safetyBufferKey + ) as? Double + let safetyBuffer = SafetyBufferPolicy.normalized(storedSafetyBuffer) + if storedSafetyBuffer != safetyBuffer { + defaults.set(safetyBuffer, forKey: Self.safetyBufferKey) + } let resetReminderCoordinator = ResetReminderCoordinator( defaults: defaults, scheduler: resetReminderScheduler @@ -112,9 +141,12 @@ final class UsageMonitor: ObservableObject { } if let data = defaults.data(forKey: Self.stateKey), let state = try? JSONDecoder().decode(StoredState.self, from: data) { - accountSnapshot = state.snapshot - samples = state.samples - legacySamplesAwaitingMigration = state.samples + let restoredSamples = state.samples.filter(\.isValid) + accountSnapshot = state.snapshot.flatMap { + $0.isValid ? $0 : nil + } + samples = restoredSamples + legacySamplesAwaitingMigration = restoredSamples previousStatus = state.previousStatus } @@ -136,7 +168,13 @@ final class UsageMonitor: ObservableObject { if defaults.object(forKey: Self.localHistoryDeletionCutoffKey) != nil { historyDeletionStatus = .pendingLocal } - recalculate() + if accountSnapshot != nil || !samples.isEmpty { + let pending = beginRecalculation() + Task { [weak self] in + guard let self else { return } + _ = await finishRecalculation(pending) + } + } if startsAutomatically { Task { [weak self] in @@ -161,24 +199,33 @@ final class UsageMonitor: ObservableObject { Timer.publish(every: 600, on: .main, in: .common) .autoconnect() .sink { [weak self] _ in - Task { @MainActor in await self?.refresh() } + Task { + @MainActor in await self?.refresh( + forceHistorySync: false + ) + } } .store(in: &cancellables) NSWorkspace.shared.notificationCenter .publisher(for: NSWorkspace.didWakeNotification) .sink { [weak self] _ in - Task { @MainActor in await self?.refresh() } + Task { + @MainActor in await self?.refresh( + forceHistorySync: false + ) + } } .store(in: &cancellables) - await refresh() + await refresh(forceHistorySync: false) } - func refresh() async { + func refresh(forceHistorySync: Bool = true) async { guard !isRefreshing else { return } isRefreshing = true defer { isRefreshing = false } + cancelLocalImport() await restoreHistoryIfAvailable() let fetchTask = Task { try await fetchUsage() } @@ -187,7 +234,9 @@ final class UsageMonitor: ObservableObject { let result = try await fetchTask.value guard let account = result.account else { historyMatchesCurrentSnapshot = false - await exchangeRestoredHistoryIfAvailable() + await exchangeRestoredHistoryIfAvailable( + force: forceHistorySync + ) accountSnapshot = result.snapshot sourceState = .available await localActivityCollector?.selectPartition( @@ -198,9 +247,13 @@ final class UsageMonitor: ObservableObject { observedAt: result.snapshot.fetchedAt, identityVerified: false ) - recalculate(now: result.snapshot.fetchedAt) + let published = await recalculate( + now: result.snapshot.fetchedAt + ) persist() - await reconcileResetReminder() + if published { + await reconcileResetReminder() + } return } let legacySamples = legacySamplesAwaitingMigration @@ -217,7 +270,7 @@ final class UsageMonitor: ObservableObject { apply(historyState) historyUsesFiles = historyState.errorMessage == nil } - let historyState = await exchangeHistory() + let historyState = await exchangeHistory(force: forceHistorySync) apply(historyState, configuredFolderName: configuredSyncDirectory?.lastPathComponent) repairInitialAccountEpochIfNeeded() let exchangeErrorMessage = historyState.errorMessage @@ -248,11 +301,15 @@ final class UsageMonitor: ObservableObject { for: newSnapshot, observedAt: newSnapshot.fetchedAt ) - recalculate(now: newSnapshot.fetchedAt) + let published = await recalculate(now: newSnapshot.fetchedAt) persist() - await reconcileResetReminder() + if published { + await reconcileResetReminder() + } } catch { - await exchangeRestoredHistoryIfAvailable() + await exchangeRestoredHistoryIfAvailable( + force: forceHistorySync + ) sourceState = .failed( (error as? CodexClientError)?.localizedDescription ?? "Couldn’t read Codex usage. Try refreshing again." @@ -267,14 +324,23 @@ final class UsageMonitor: ObservableObject { identityVerified: false ) } - recalculate() + _ = await recalculate() persist() } } func updateSafetyBuffer(_ value: Double) { - recalculate(safetyBuffer: value) - persist() + let value = SafetyBufferPolicy.normalized(value) + defaults.set(value, forKey: Self.safetyBufferKey) + let pending = beginRecalculation(safetyBuffer: value) + Task { [weak self] in + guard let self else { return } + let published = await finishRecalculation(pending) + persist() + if published { + await reconcileResetReminder() + } + } } func setResetReminderEnabled(_ isEnabled: Bool) async { @@ -359,6 +425,7 @@ final class UsageMonitor: ObservableObject { planType: String? = nil, observedAt: Date ) async { + cancelLocalImport() let partition: AccountHistoryPartition let authState: String let previousAuthState = defaults.string(forKey: Self.historyAuthStateKey) @@ -425,7 +492,7 @@ final class UsageMonitor: ObservableObject { localActivityCollection = .unavailable( "Codex local records are unavailable" ) - recalculate() + _ = await recalculate() if historyPrepared { persist() } @@ -475,6 +542,7 @@ final class UsageMonitor: ObservableObject { guard !isUpdatingHistory else { return } isUpdatingHistory = true defer { isUpdatingHistory = false } + cancelLocalImport() let localDeletedAt = Date() NotificationCenter.default.post( name: .codexAssistedHistoryDeleted, @@ -528,7 +596,7 @@ final class UsageMonitor: ObservableObject { accountFacts: snapshot.accountFacts ) } - recalculate() + _ = await recalculate() persist() } @@ -536,6 +604,7 @@ final class UsageMonitor: ObservableObject { guard !isUpdatingHistory else { return } isUpdatingHistory = true defer { isUpdatingHistory = false } + cancelLocalImport() let assistedDeletionCutoff = defaults.object( forKey: Self.localHistoryDeletionCutoffKey ) as? Date @@ -612,7 +681,7 @@ final class UsageMonitor: ObservableObject { } else { false } - recalculate() + _ = await recalculate() persist() } @@ -625,6 +694,7 @@ final class UsageMonitor: ObservableObject { guard !isUpdatingHistory, canRebuildAvailableHistory else { return } isUpdatingHistory = true defer { isUpdatingHistory = false } + cancelLocalImport() do { let result = try await fetchUsage() guard let account = result.account else { @@ -656,16 +726,18 @@ final class UsageMonitor: ObservableObject { for: snapshot, observedAt: snapshot.fetchedAt ) - recalculate(now: snapshot.fetchedAt) + let published = await recalculate(now: snapshot.fetchedAt) persist() - await reconcileResetReminder() + if published { + await reconcileResetReminder() + } } catch let error as CodexClientError { sourceState = .failed(error.localizedDescription) - recalculate() + _ = await recalculate() persist() } catch { sourceState = .failed("Couldn’t read Codex usage. Try again.") - recalculate() + _ = await recalculate() persist() } } @@ -673,36 +745,106 @@ final class UsageMonitor: ObservableObject { private func recalculate( safetyBuffer: Double? = nil, now: Date = Date() + ) async -> Bool { + await finishRecalculation( + beginRecalculation(safetyBuffer: safetyBuffer, now: now) + ) + } + + private func beginRecalculation( + safetyBuffer: Double? = nil, + now: Date = Date(), + analyticsExploration: AnalyticsExplorationState? = nil, + insightDispositions: [String: InsightDisposition]? = nil + ) -> ( + generation: UInt64, + task: Task ) { + let input = evaluationInput( + safetyBuffer: safetyBuffer, + now: now, + analyticsExploration: analyticsExploration, + insightDispositions: insightDispositions + ) + evaluationGeneration &+= 1 + let generation = evaluationGeneration + let previousTask = evaluationTask + previousTask?.cancel() + let evaluateUsage = evaluateUsage + let task: Task = Task.detached( + priority: .userInitiated + ) { + _ = await previousTask?.value + guard !Task.isCancelled else { return nil } + let snapshot = evaluateUsage(input) + return Task.isCancelled ? nil : snapshot + } + evaluationTask = task + return (generation, task) + } + + private func evaluationInput( + safetyBuffer: Double? = nil, + now: Date = Date(), + analyticsExploration: AnalyticsExplorationState? = nil, + insightDispositions: [String: InsightDisposition]? = nil + ) -> UsageIntelligenceInput { let storedBuffer = defaults.object(forKey: Self.safetyBufferKey) as? Double - let buffer = safetyBuffer ?? storedBuffer ?? 3 - readerSnapshot = UsageIntelligenceEngine.evaluate( - UsageIntelligenceInput( - account: accountSnapshot, - samples: historyMatchesCurrentSnapshot ? samples : [], - safetyBuffer: buffer, - sourceState: sourceState, - now: now, - previousStatus: previousStatus, - accountPartitionID: historyPartition.id, - accountEpochStartedAt: accountEpochStartedAt, - localActivityFacts: localActivityCollection.facts, - localActivityHistoryFacts: - localActivityCollection.facts, - localActivityObservation: localActivityCollection.observation, - localTaskProjections: localActivityCollection.projections, - analyticsExploration: - AnalyticsWorkspaceStore.restoredState( - from: defaults - ), - insightDispositions: - AnalyticsWorkspaceStore + let buffer = SafetyBufferPolicy.normalized( + safetyBuffer ?? storedBuffer + ) + return UsageIntelligenceInput( + account: accountSnapshot, + samples: historyMatchesCurrentSnapshot ? samples : [], + safetyBuffer: buffer, + sourceState: sourceState, + now: now, + previousStatus: previousStatus, + accountPartitionID: historyPartition.id, + accountEpochStartedAt: accountEpochStartedAt, + localActivityFacts: localActivityCollection.facts, + localActivityHistoryFacts: localActivityCollection.facts, + localActivityObservation: localActivityCollection.observation, + localTaskProjections: localActivityCollection.projections, + localActivityContentRevision: + localActivityCollection.contentRevision, + reusableLocalAggregates: + readerSnapshot.reusableLocalAggregates, + analyticsExploration: + analyticsExploration + ?? AnalyticsWorkspaceStore.restoredState(from: defaults), + insightDispositions: + insightDispositions + ?? AnalyticsWorkspaceStore .restoredInsightDispositions(from: defaults) - ) ) - if let status = readerSnapshot.guidance?.status { + } + + private func finishRecalculation( + _ pending: ( + generation: UInt64, + task: Task + ) + ) async -> Bool { + let snapshot = await withTaskCancellationHandler { + await pending.task.value + } onCancel: { + pending.task.cancel() + } + if pending.generation == evaluationGeneration { + evaluationTask = nil + } + guard let snapshot, + !Task.isCancelled, + pending.generation == evaluationGeneration, + !pending.task.isCancelled else { + return false + } + readerSnapshot = snapshot + if let status = snapshot.guidance?.status { previousStatus = status } + return true } func analyticsPreferencesDidChange( @@ -717,6 +859,17 @@ final class UsageMonitor: ObservableObject { input, dispositions: dispositions ) + guard evaluationTask != nil else { return } + let pending = beginRecalculation( + analyticsExploration: exploration, + insightDispositions: dispositions + ) + Task { [weak self] in + guard let self else { return } + if await finishRecalculation(pending) { + await reconcileResetReminder() + } + } } private func reconcileResetReminder() async { @@ -752,6 +905,7 @@ final class UsageMonitor: ObservableObject { observedAt: Date, identityVerified: Bool = true ) async { + cancelLocalImport() guard let interval = UsageIntelligenceEngine.tokenActivityInterval( account: snapshot, samples: historyMatchesCurrentSnapshot ? samples : [], @@ -768,6 +922,9 @@ final class UsageMonitor: ObservableObject { ) return } + localActivityCollection = .unavailable( + "Codex local records are unavailable" + ) let collection = await localActivityCollector.refresh( interval: interval, observedAt: observedAt @@ -780,6 +937,69 @@ final class UsageMonitor: ObservableObject { : collection.loweringCoverage( "Codex account identity could not be checked" ) + guard await localActivityCollector.hasPendingImport() else { return } + let generation = localImportGeneration + localImportTask = Task(priority: .background) { [weak self] in + try? await Task.sleep(for: .milliseconds(500)) + await self?.continueLocalActivityImport( + with: localActivityCollector, + interval: interval, + observedAt: observedAt, + identityVerified: identityVerified, + generation: generation + ) + } + } + + private func continueLocalActivityImport( + with collector: LocalActivityCollector, + interval: DateInterval, + observedAt: Date, + identityVerified: Bool, + generation: UInt64 + ) async { + var latest: LocalActivityCollection? + while !Task.isCancelled, await collector.hasPendingImport() { + while isRefreshing, !Task.isCancelled { + try? await Task.sleep(for: .milliseconds(100)) + } + guard generation == localImportGeneration else { return } + latest = nil + let collection = await collector.refresh( + interval: interval, + observedAt: observedAt, + refreshMetadata: false + ) + guard !Task.isCancelled, + generation == localImportGeneration else { + return + } + latest = collection + if await collector.hasPendingImport() { + try? await Task.sleep(for: .milliseconds(250)) + } + } + guard let latest, + !Task.isCancelled, + generation == localImportGeneration else { + return + } + localActivityCollection = identityVerified + ? latest + : latest.loweringCoverage( + "Codex account identity could not be checked" + ) + _ = await recalculate(now: observedAt) + persist() + if generation == localImportGeneration { + localImportTask = nil + } + } + + private func cancelLocalImport() { + localImportGeneration &+= 1 + localImportTask?.cancel() + localImportTask = nil } private func persist() { @@ -889,8 +1109,13 @@ final class UsageMonitor: ObservableObject { } } - private func exchangeHistory() async -> UsageHistory.State { + private func exchangeHistory(force: Bool) async -> UsageHistory.State { if let configuredSyncDirectory { + if await history.isConnected(to: configuredSyncDirectory) { + return force + ? await history.synchronize() + : await history.synchronizeIfDue() + } let state = await history.connect( to: configuredSyncDirectory, accountIdentity: historyAccountIdentity, @@ -901,13 +1126,15 @@ final class UsageMonitor: ObservableObject { historyConnectionActive = state.folderName != nil return state } - return await history.synchronize() + return force + ? await history.synchronize() + : await history.synchronizeIfDue() } - private func exchangeRestoredHistoryIfAvailable() async { + private func exchangeRestoredHistoryIfAvailable(force: Bool) async { guard restoredFileStoreAvailable else { return } await prepareHistory(legacySamples: []) - let state = await exchangeHistory() + let state = await exchangeHistory(force: force) apply(state, configuredFolderName: configuredSyncDirectory?.lastPathComponent) } diff --git a/Sources/CodexLimits/UsagePerToken.swift b/Sources/CodexLimits/UsagePerToken.swift index b793d51..b85db19 100644 --- a/Sources/CodexLimits/UsagePerToken.swift +++ b/Sources/CodexLimits/UsagePerToken.swift @@ -238,10 +238,11 @@ enum WeeklyUsageEvidenceBuilder { accountPartitionID: String, limitID: String, currentReset: Date, - compatibleTokenSources: Set + compatibleTokenSources: Set, + factIndex: LocalActivityFactIndex? = nil ) -> WeeklyUsageEvidenceSet { let grouped = Dictionary(grouping: samples, by: \.resetsAt) - let localFactIndex = LocalActivityFactIndex(localFacts) + let localFactIndex = factIndex ?? LocalActivityFactIndex(localFacts) let latestComparisonBreak = samples .lazy .filter(\.comparisonBreak) @@ -251,6 +252,7 @@ enum WeeklyUsageEvidenceBuilder { buildInterval( samples: samples, localFactIndex: localFactIndex, + localFacts: localFacts, localObservation: localObservation, accountPartitionID: accountPartitionID, limitID: limitID, @@ -274,6 +276,7 @@ enum WeeklyUsageEvidenceBuilder { private static func buildInterval( samples: [UsageSample], localFactIndex: LocalActivityFactIndex, + localFacts: [LocalActivityFact], localObservation: LocalActivityObservation, accountPartitionID: String, limitID: String, @@ -380,11 +383,15 @@ enum WeeklyUsageEvidenceBuilder { boundedTokenReadings, boundedTokenReadings.dropFirst() ).contains { $0.1.1 < $0.0.1 } - let tokenDifference = endTokens.subtractingReportingOverflow( - startTokens + let (accountTokens, tokenDifferenceOverflowed) = + endTokens.subtractingReportingOverflow( + startTokens + ) + guard !tokenDifferenceOverflowed else { return nil } + let intervalFacts = localFactIndex.facts( + in: interval, + from: localFacts ) - guard !tokenDifference.overflow else { return nil } - let intervalFacts = localFactIndex.facts(in: interval) let workload = workload( facts: intervalFacts, interval: interval @@ -394,7 +401,6 @@ enum WeeklyUsageEvidenceBuilder { interval: interval, intervalBreakReason: workload.sourceBreakReason ) - let accountTokens = tokenDifference.partialValue let tokenDefinitionsAlign = tokenDefinitionsAlign( facts: intervalFacts, compatibleSources: compatibleTokenSources @@ -582,8 +588,8 @@ enum WeeklyUsageEvidenceBuilder { reasoning[level] = reasoningTokens } guard let tokenDelta, - tokenDelta.observedComponents.contains(.input), - tokenDelta.observedComponents.contains(.cachedInput) else { + tokenDelta.observes(.input), + tokenDelta.observes(.cachedInput) else { hasCacheEvidence = false continue } diff --git a/Sources/CodexLimits/UsageReceipts.swift b/Sources/CodexLimits/UsageReceipts.swift index 8af9803..525bc1a 100644 --- a/Sources/CodexLimits/UsageReceipts.swift +++ b/Sources/CodexLimits/UsageReceipts.swift @@ -201,6 +201,47 @@ struct UsageReceiptSlice: Equatable, Sendable { let receiptReason: String? } +struct UsageReceiptSummary: Equatable, Identifiable, Sendable { + let rootTaskID: String + let projectLabel: String? + let tokens: Int64 + let taskCount: Int + let coverage: CoverageLevel + let reason: String? + + var id: String { rootTaskID } + + var displayTaskID: String { + String(rootTaskID.prefix(8)) + } + + var accessibilityValue: String { + var parts = [ + "Task \(displayTaskID)", + "\(tokens) local tokens", + "\(taskCount) \(taskCount == 1 ? "task" : "tasks") in the tree" + ] + if let projectLabel { + parts.insert("Project \(projectLabel)", at: 0) + } + parts.append("\(coverage.displayName) coverage") + if let reason { + parts.append(reason) + } + return parts.joined(separator: ", ") + } +} + +struct UsageReceiptOverview: Equatable, Sendable { + let receipts: [UsageReceiptSummary] + let totalTokens: Int64 + let unattributedTokens: Int64 + let coverage: CoverageLevel + let reason: String? + let receiptCoverage: CoverageLevel + let receiptReason: String? +} + struct UsageReceiptFilterOptions: Equatable, Sendable { let projects: [String] let taskTrees: [String] @@ -216,34 +257,278 @@ struct UsageReceiptSnapshot: Equatable, Sendable { fileprivate let observation: LocalActivityObservation let interval: DateInterval - func slice( + func updating( + interval: DateInterval, + observation: LocalActivityObservation + ) -> UsageReceiptSnapshot { + UsageReceiptSnapshot( + contributions: contributions, + diagnostics: diagnostics, + projections: projections, + taskIDsByRoot: taskIDsByRoot, + observation: observation, + interval: interval + ) + } + + func overview( in selectedInterval: DateInterval, filters: WorkspaceFilters - ) -> UsageReceiptSlice { - let inRange = contributions.filter { - Self.contains($0.date, in: selectedInterval) + ) -> UsageReceiptOverview { + var accumulators: [String: SummaryAccumulator] = [:] + var totalTokens: Int64 = 0 + var unattributedTokens: Int64 = 0 + var totalOverflow = false + var unattributedOverflow = false + var hasUnboundedCounter = false + var hasMissingProject = false + var hasTokenContribution = false + var contributionGaps = FilterGaps() + var diagnosticGaps = FilterGaps() + + for contribution in contributions + where Self.contains(contribution.date, in: selectedInterval) { + contributionGaps.observe( + Self.metadata(for: contribution), + filters: filters + ) + guard Self.matches(contribution, filters: filters) else { + continue + } + hasUnboundedCounter = + hasUnboundedCounter || contribution.hasUnboundedCounter + guard !totalOverflow else { continue } + let total = totalTokens.addingReportingOverflow( + contribution.tokens + ) + totalOverflow = total.overflow + totalTokens = total.partialValue + if let rootTaskID = contribution.rootTaskID { + var accumulator = accumulators[rootTaskID] + ?? SummaryAccumulator() + accumulator.hasUnboundedCounter = + accumulator.hasUnboundedCounter + || contribution.hasUnboundedCounter + if contribution.tokens > 0 { + accumulator.add(contribution) + } + accumulators[rootTaskID] = accumulator + } + guard contribution.tokens > 0 else { continue } + hasTokenContribution = true + hasMissingProject = + hasMissingProject || contribution.projectLabel == nil + guard contribution.rootTaskID == nil else { + continue + } + let unattributed = unattributedTokens.addingReportingOverflow( + contribution.tokens + ) + unattributedOverflow = + unattributedOverflow || unattributed.overflow + unattributedTokens = unattributed.partialValue } - let diagnosticsInRange = diagnostics.filter { - $0.intersects(selectedInterval) + + for diagnostic in diagnostics + where diagnostic.intersects(selectedInterval) { + diagnosticGaps.observe( + Self.metadata(for: diagnostic), + filters: filters + ) + guard Self.matches(diagnostic, filters: filters), + let rootTaskID = diagnostic.rootTaskID else { + continue + } + var accumulator = accumulators[rootTaskID] + ?? SummaryAccumulator() + accumulator.add(diagnostic) + accumulators[rootTaskID] = accumulator } - let filterGapReason = Self.filterGapReason( - in: inRange, - filters: filters - ) ?? Self.filterGapReason( - in: diagnosticsInRange, + + guard !totalOverflow, !unattributedOverflow, + !accumulators.values.contains(where: \.tokensOverflow) else { + return UsageReceiptOverview( + receipts: [], + totalTokens: 0, + unattributedTokens: 0, + coverage: .unavailable, + reason: "Local token total is invalid", + receiptCoverage: .unavailable, + receiptReason: "Local token total is invalid" + ) + } + + let filterGapReason = contributionGaps.reason( + subject: "Some local activity", + verb: "has" + ) ?? diagnosticGaps.reason( + subject: "Some local diagnostics", + verb: "have" + ) + let summaries: [UsageReceiptSummary] = accumulators.compactMap { + rootTaskID, accumulator -> UsageReceiptSummary? in + guard accumulator.hasReceiptEvidence else { return nil } + let evidence = receiptEvidence( + hasMissingProject: accumulator.hasMissingProject, + hasUnboundedCounter: accumulator.hasUnboundedCounter, + filterGapReason: filterGapReason + ) + return UsageReceiptSummary( + rootTaskID: rootTaskID, + projectLabel: accumulator.projectLabel, + tokens: accumulator.tokens, + taskCount: max( + taskIDsByRoot[rootTaskID]?.count ?? 0, + 1 + ), + coverage: evidence.0, + reason: evidence.1 + ) + } + .sorted { + ($0.projectLabel ?? "", $0.rootTaskID) + < ($1.projectLabel ?? "", $1.rootTaskID) + } + let coverage = sliceEvidence( + hasTokenContribution: hasTokenContribution, + hasUnboundedCounter: hasUnboundedCounter, + hasMissingProject: hasMissingProject, + hasUnattributedTokens: unattributedTokens > 0, + hasReceipts: !summaries.isEmpty, + filterGapReason: filterGapReason + ) + let combinedReceiptEvidence = Self.receiptEvidence(summaries) + return UsageReceiptOverview( + receipts: summaries, + totalTokens: totalTokens, + unattributedTokens: unattributedTokens, + coverage: coverage.0, + reason: coverage.1, + receiptCoverage: combinedReceiptEvidence.0, + receiptReason: combinedReceiptEvidence.1 + ) + } + + func receipt( + rootTaskID: String, + in selectedInterval: DateInterval, + filters: WorkspaceFilters + ) -> UsageReceipt? { + return slice( + in: selectedInterval, + filters: filters, + onlyRootTaskID: rootTaskID + ).receipts.first + } + + func localTokenSlice( + in selectedInterval: DateInterval, + filters: WorkspaceFilters + ) -> LocalTokenActivitySlice { + let overview = overview( + in: selectedInterval, filters: filters ) - let selected = inRange.filter { - Self.matches($0, filters: filters) + let selected = contributions.filter { + Self.contains($0.date, in: selectedInterval) + && $0.tokens > 0 + && Self.matches($0, filters: filters) } - let selectedDiagnostics = diagnosticsInRange.filter { - Self.matches($0, filters: filters) + return LocalTokenActivitySlice( + tokens: overview.totalTokens, + points: cumulativePoints(selected), + coverage: overview.coverage, + reason: overview.reason + ) + } + + func slice( + in selectedInterval: DateInterval, + filters: WorkspaceFilters + ) -> UsageReceiptSlice { + slice( + in: selectedInterval, + filters: filters, + onlyRootTaskID: nil + ) + } + + private func slice( + in selectedInterval: DateInterval, + filters: WorkspaceFilters, + onlyRootTaskID: String? + ) -> UsageReceiptSlice { + var groupedContributions: [String: [Contribution]] = [:] + var groupedDiagnostics: [String: [DiagnosticContribution]] = [:] + var tokenContributions: [Contribution] = [] + var unboundedRoots = Set() + var totalTokens: Int64 = 0 + var unattributedTokens: Int64 = 0 + var totalOverflow = false + var unattributedOverflow = false + var hasUnboundedCounter = false + var hasMissingProject = false + var contributionGaps = FilterGaps() + var diagnosticGaps = FilterGaps() + + for contribution in contributions + where Self.contains(contribution.date, in: selectedInterval) { + contributionGaps.observe( + Self.metadata(for: contribution), + filters: filters + ) + guard Self.matches(contribution, filters: filters), + onlyRootTaskID == nil + || contribution.rootTaskID == onlyRootTaskID else { + continue + } + let total = totalTokens.addingReportingOverflow( + contribution.tokens + ) + totalOverflow = totalOverflow || total.overflow + totalTokens = total.partialValue + hasUnboundedCounter = + hasUnboundedCounter || contribution.hasUnboundedCounter + if contribution.hasUnboundedCounter, + let rootTaskID = contribution.rootTaskID { + unboundedRoots.insert(rootTaskID) + } + guard contribution.tokens > 0 else { continue } + tokenContributions.append(contribution) + hasMissingProject = + hasMissingProject || contribution.projectLabel == nil + if let rootTaskID = contribution.rootTaskID { + groupedContributions[rootTaskID, default: []] + .append(contribution) + } else { + let unattributed = unattributedTokens.addingReportingOverflow( + contribution.tokens + ) + unattributedOverflow = + unattributedOverflow || unattributed.overflow + unattributedTokens = unattributed.partialValue + } } - let tokenContributions = selected.filter { $0.tokens > 0 } - let hasUnboundedCounter = selected.contains { - $0.hasUnboundedCounter + + for diagnostic in diagnostics + where diagnostic.intersects(selectedInterval) { + diagnosticGaps.observe( + Self.metadata(for: diagnostic), + filters: filters + ) + guard Self.matches(diagnostic, filters: filters), + onlyRootTaskID == nil + || diagnostic.rootTaskID == onlyRootTaskID, + let rootTaskID = diagnostic.rootTaskID else { + continue + } + groupedDiagnostics[rootTaskID, default: []].append(diagnostic) } - guard let total = Self.sum(selected.map(\.tokens)) else { + + guard !totalOverflow, !unattributedOverflow, + groupedContributions.values.allSatisfy({ + Self.sum($0.lazy.map(\.tokens)) != nil + }) else { return UsageReceiptSlice( receipts: [], totalTokens: 0, @@ -255,25 +540,18 @@ struct UsageReceiptSnapshot: Equatable, Sendable { receiptReason: "Local token total is invalid" ) } - let groupedContributions = Dictionary( - grouping: tokenContributions.compactMap { contribution in - contribution.rootTaskID.map { ($0, contribution) } - }, - by: \.0 - ) - let groupedDiagnostics = Dictionary( - grouping: selectedDiagnostics.compactMap { diagnostic in - diagnostic.rootTaskID.map { ($0, diagnostic) } - }, - by: \.0 + let filterGapReason = contributionGaps.reason( + subject: "Some local activity", + verb: "has" + ) ?? diagnosticGaps.reason( + subject: "Some local diagnostics", + verb: "have" ) let rootTaskIDs = Set(groupedContributions.keys) .union(groupedDiagnostics.keys) let receipts = rootTaskIDs.map { rootTaskID in - let contributions = (groupedContributions[rootTaskID] ?? []) - .map(\.1) - let receiptDiagnostics = (groupedDiagnostics[rootTaskID] ?? []) - .map(\.1) + let contributions = groupedContributions[rootTaskID] ?? [] + let receiptDiagnostics = groupedDiagnostics[rootTaskID] ?? [] let project = contributions.compactMap(\.projectLabel).first ?? receiptDiagnostics.compactMap(\.projectLabel).first let hasMissingProject = contributions.contains { @@ -281,32 +559,11 @@ struct UsageReceiptSnapshot: Equatable, Sendable { } || receiptDiagnostics.contains { $0.projectLabel == nil } - let receiptHasUnboundedCounter = selected.contains { - $0.rootTaskID == rootTaskID && $0.hasUnboundedCounter - } - let receiptEvidence: (CoverageLevel, String) - switch observation { - case let .unavailable(message): - receiptEvidence = (.unavailable, message) - case let .gap(_, _, message): - receiptEvidence = (.low, message) - case .continuous: - if receiptHasUnboundedCounter { - receiptEvidence = ( - .low, - "Local token activity starts from an unbounded counter" - ) - } else if let filterGapReason { - receiptEvidence = (.partial, filterGapReason) - } else { - receiptEvidence = hasMissingProject - ? (.partial, "Project metadata is missing") - : ( - .partial, - "Task Tree may omit Review and Guardian Tasks" - ) - } - } + let receiptEvidence = receiptEvidence( + hasMissingProject: hasMissingProject, + hasUnboundedCounter: unboundedRoots.contains(rootTaskID), + filterGapReason: filterGapReason + ) let taskTree = Self.taskTree( rootTaskID: rootTaskID, contributions: contributions, @@ -321,7 +578,9 @@ struct UsageReceiptSnapshot: Equatable, Sendable { return UsageReceipt( rootTaskID: rootTaskID, projectLabel: project, - tokens: contributions.reduce(0) { $0 + $1.tokens }, + tokens: Self.sum( + contributions.lazy.map(\.tokens) + ) ?? 0, interval: selectedInterval, taskCount: taskTree.taskCount, taskTree: taskTree, @@ -332,6 +591,7 @@ struct UsageReceiptSnapshot: Equatable, Sendable { $0.context?.reasoning }, diagnostics: Self.diagnosticSummary( + contributions, receiptDiagnostics, selectedInterval: selectedInterval, observation: observation @@ -344,48 +604,22 @@ struct UsageReceiptSnapshot: Equatable, Sendable { ($0.projectLabel ?? "", $0.rootTaskID) < ($1.projectLabel ?? "", $1.rootTaskID) } - let unattributed = tokenContributions - .filter { $0.rootTaskID == nil } - .reduce(0) { $0 + $1.tokens } - let coverage: CoverageLevel - let reason: String? - switch observation { - case .unavailable(let message): - coverage = .unavailable - reason = message - case .gap(_, _, let message): - coverage = .low - reason = message - case .continuous: - if hasUnboundedCounter { - coverage = .low - reason = "Local token activity starts from an unbounded counter" - } else if let filterGapReason { - coverage = tokenContributions.isEmpty ? .low : .partial - reason = filterGapReason - } else if tokenContributions.isEmpty { - coverage = .notApplicable - reason = "No local token activity was observed" - } else if unattributed > 0 || tokenContributions.contains( - where: { $0.projectLabel == nil } - ) { - coverage = receipts.isEmpty ? .low : .partial - reason = unattributed > 0 - ? "Task metadata is missing" - : "Project metadata is missing" - } else { - coverage = .high - reason = "Only local activity on this Mac is observed" - } - } + let coverage = sliceEvidence( + hasTokenContribution: !tokenContributions.isEmpty, + hasUnboundedCounter: hasUnboundedCounter, + hasMissingProject: hasMissingProject, + hasUnattributedTokens: unattributedTokens > 0, + hasReceipts: !receipts.isEmpty, + filterGapReason: filterGapReason + ) let combinedReceiptEvidence = Self.receiptEvidence(receipts) return UsageReceiptSlice( receipts: receipts, - totalTokens: total, - unattributedTokens: unattributed, + totalTokens: totalTokens, + unattributedTokens: unattributedTokens, points: cumulativePoints(tokenContributions), - coverage: coverage, - reason: reason, + coverage: coverage.0, + reason: coverage.1, receiptCoverage: combinedReceiptEvidence.0, receiptReason: combinedReceiptEvidence.1 ) @@ -394,37 +628,46 @@ struct UsageReceiptSnapshot: Equatable, Sendable { func filterOptions( in selectedInterval: DateInterval ) -> UsageReceiptFilterOptions { - let selected = contributions.filter { - Self.contains($0.date, in: selectedInterval) && $0.tokens > 0 + var projects = Set() + var taskTrees = Set() + var models = Set() + var reasoningLevels = Set() + for contribution in contributions + where Self.contains(contribution.date, in: selectedInterval) + && contribution.tokens > 0 { + if let project = contribution.projectLabel { + projects.insert(project) + } + if let taskTree = contribution.rootTaskID { + taskTrees.insert(taskTree) + } + if let model = contribution.context?.effectiveModel { + models.insert(model) + } + if let reasoning = contribution.context?.reasoning { + reasoningLevels.insert(reasoning) + } } - let selectedDiagnostics = diagnostics.filter { - $0.intersects(selectedInterval) + for diagnostic in diagnostics + where diagnostic.intersects(selectedInterval) { + if let project = diagnostic.projectLabel { + projects.insert(project) + } + if let taskTree = diagnostic.rootTaskID { + taskTrees.insert(taskTree) + } + if let model = diagnostic.context?.effectiveModel { + models.insert(model) + } + if let reasoning = diagnostic.context?.reasoning { + reasoningLevels.insert(reasoning) + } } return UsageReceiptFilterOptions( - projects: Set(selected.compactMap(\.projectLabel)) - .union(selectedDiagnostics.compactMap(\.projectLabel)) - .sorted(), - taskTrees: Set(selected.compactMap(\.rootTaskID)) - .union(selectedDiagnostics.compactMap(\.rootTaskID)) - .sorted(), - models: Set( - selected.compactMap { $0.context?.effectiveModel } - ) - .union( - selectedDiagnostics.compactMap { - $0.context?.effectiveModel - } - ) - .sorted(), - reasoningLevels: Set( - selected.compactMap { $0.context?.reasoning } - ) - .union( - selectedDiagnostics.compactMap { - $0.context?.reasoning - } - ) - .sorted() + projects: projects.sorted(), + taskTrees: taskTrees.sorted(), + models: models.sorted(), + reasoningLevels: reasoningLevels.sorted() ) } @@ -433,24 +676,27 @@ struct UsageReceiptSnapshot: Equatable, Sendable { let tokens: Int64 let rootTaskID: String? let projectLabel: String? - let context: LocalActivityContext? + let sharedContext: SharedContext? let tokenSource: LocalActivitySourceKind let hasUnboundedCounter: Bool + let tokenDelta: LocalTokenUsage? + let contextUsage: LocalTokenUsage? + + var context: LocalActivityContext? { sharedContext?.value } } fileprivate struct DiagnosticContribution: Equatable, Sendable { let date: Date let rootTaskID: String? let projectLabel: String? - let context: LocalActivityContext? - let key: LocalActivityFactKey - let value: LocalActivityFactValue? - let tokenDelta: LocalTokenUsage? - let availability: LocalActivityAvailability - let reason: String? + let sharedContext: SharedContext? + let payload: DiagnosticPayload let eventID: String let source: LocalActivitySourceKind + var context: LocalActivityContext? { sharedContext?.value } + var key: LocalActivityFactKey { payload.key } + var value: LocalActivityFactValue? { payload.value } func intersects(_ interval: DateInterval) -> Bool { if let durationInterval { return durationInterval.start < interval.end @@ -465,7 +711,7 @@ struct UsageReceiptSnapshot: Equatable, Sendable { } var durationInterval: DateInterval? { - guard case let .duration(duration) = value, + guard let duration = payload.duration, duration.completedAt > duration.startedAt else { return nil } @@ -476,6 +722,56 @@ struct UsageReceiptSnapshot: Equatable, Sendable { } } + fileprivate final class SharedContext: Sendable, Equatable { + let value: LocalActivityContext + + init(_ value: LocalActivityContext) { + self.value = value + } + + static func == (lhs: SharedContext, rhs: SharedContext) -> Bool { + lhs.value == rhs.value + } + } + + fileprivate enum DiagnosticPayload: Equatable, Sendable { + case context(LocalTokenUsage?) + case time(LocalTurnTiming?) + case tool(String?) + case duration(LocalActivityFactKey, LocalActivityDuration?) + case compaction + + var key: LocalActivityFactKey { + switch self { + case .context: .context + case .time: .time + case .tool: .tool + case let .duration(key, _): key + case .compaction: .compaction + } + } + + var value: LocalActivityFactValue? { + switch self { + case let .context(value): + value.map(LocalActivityFactValue.tokens) + case let .time(value): + value.map(LocalActivityFactValue.turnTiming) + case let .tool(value): + value.map(LocalActivityFactValue.text) + case let .duration(_, value): + value.map(LocalActivityFactValue.duration) + case .compaction: + .count(1) + } + } + + var duration: LocalActivityDuration? { + guard case let .duration(_, value) = self else { return nil } + return value + } + } + private struct FilterMetadata { let projectID: String? let taskTreeID: String? @@ -490,6 +786,193 @@ struct UsageReceiptSnapshot: Equatable, Sendable { case reasoning } + private struct FilterGaps { + var project = false + var taskTree = false + var model = false + var reasoning = false + + mutating func observe( + _ metadata: FilterMetadata, + filters: WorkspaceFilters + ) { + if filters.projectID != nil, + metadata.projectID == nil, + UsageReceiptSnapshot.couldMatch( + metadata, + filters: filters, + ignoring: .project + ) { + project = true + } + if filters.taskTreeID != nil, + metadata.taskTreeID == nil, + UsageReceiptSnapshot.couldMatch( + metadata, + filters: filters, + ignoring: .taskTree + ) { + taskTree = true + } + if filters.model != nil, + metadata.model == nil, + UsageReceiptSnapshot.couldMatch( + metadata, + filters: filters, + ignoring: .model + ) { + model = true + } + if filters.reasoning != nil, + metadata.reasoning == nil, + UsageReceiptSnapshot.couldMatch( + metadata, + filters: filters, + ignoring: .reasoning + ) { + reasoning = true + } + } + + func reason(subject: String, verb: String) -> String? { + if project { + return "\(subject) \(verb) no Project metadata" + } + if taskTree { + return "\(subject) \(verb) no Task metadata" + } + if model { + return "\(subject) \(verb) no model metadata" + } + if reasoning { + return "\(subject) \(verb) no reasoning metadata" + } + return nil + } + } + + private struct SummaryAccumulator { + var projectLabel: String? + var tokens: Int64 = 0 + var tokensOverflow = false + var hasMissingProject = false + var hasUnboundedCounter = false + var hasReceiptEvidence = false + + mutating func add(_ contribution: Contribution) { + hasReceiptEvidence = true + projectLabel = projectLabel ?? contribution.projectLabel + hasMissingProject = + hasMissingProject || contribution.projectLabel == nil + let result = tokens.addingReportingOverflow(contribution.tokens) + tokensOverflow = tokensOverflow || result.overflow + tokens = result.partialValue + } + + mutating func add(_ diagnostic: DiagnosticContribution) { + hasReceiptEvidence = true + projectLabel = projectLabel ?? diagnostic.projectLabel + hasMissingProject = + hasMissingProject || diagnostic.projectLabel == nil + } + } + + private static func metadata( + for contribution: Contribution + ) -> FilterMetadata { + FilterMetadata( + projectID: contribution.projectLabel, + taskTreeID: contribution.rootTaskID, + model: contribution.context?.effectiveModel, + reasoning: contribution.context?.reasoning + ) + } + + private static func metadata( + for diagnostic: DiagnosticContribution + ) -> FilterMetadata { + FilterMetadata( + projectID: diagnostic.projectLabel, + taskTreeID: diagnostic.rootTaskID, + model: diagnostic.context?.effectiveModel, + reasoning: diagnostic.context?.reasoning + ) + } + + private func receiptEvidence( + hasMissingProject: Bool, + hasUnboundedCounter: Bool, + filterGapReason: String? + ) -> (CoverageLevel, String) { + switch observation { + case let .unavailable(message): + return (.unavailable, message) + case let .gap(_, _, message): + return (.low, message) + case .continuous: + if hasUnboundedCounter { + return ( + .low, + "Local token activity starts from an unbounded counter" + ) + } + if let filterGapReason { + return (.partial, filterGapReason) + } + if hasMissingProject { + return (.partial, "Project metadata is missing") + } + return ( + .partial, + "Task Tree may omit Review and Guardian Tasks" + ) + } + } + + private func sliceEvidence( + hasTokenContribution: Bool, + hasUnboundedCounter: Bool, + hasMissingProject: Bool, + hasUnattributedTokens: Bool, + hasReceipts: Bool, + filterGapReason: String? + ) -> (CoverageLevel, String?) { + switch observation { + case let .unavailable(message): + return (.unavailable, message) + case let .gap(_, _, message): + return (.low, message) + case .continuous: + if hasUnboundedCounter { + return ( + .low, + "Local token activity starts from an unbounded counter" + ) + } + if let filterGapReason { + return ( + hasTokenContribution ? .partial : .low, + filterGapReason + ) + } + if !hasTokenContribution { + return ( + .notApplicable, + "No local token activity was observed" + ) + } + if hasUnattributedTokens || hasMissingProject { + return ( + hasReceipts ? .partial : .low, + hasUnattributedTokens + ? "Task metadata is missing" + : "Project metadata is missing" + ) + } + return (.high, "Only local activity on this Mac is observed") + } + } + private static func taskTree( rootTaskID: String, contributions: [Contribution], @@ -599,6 +1082,7 @@ struct UsageReceiptSnapshot: Equatable, Sendable { Set(turnContributions.map(\.tokenSource)) ).sorted { $0.rawValue < $1.rawValue }, diagnostics: diagnosticSummary( + turnContributions, diagnostics, selectedInterval: selectedInterval, observation: observation @@ -878,17 +1362,27 @@ struct UsageReceiptSnapshot: Equatable, Sendable { } private static func diagnosticSummary( + _ contributions: [Contribution], _ diagnostics: [DiagnosticContribution], selectedInterval: DateInterval, observation: LocalActivityObservation ) -> UsageReceiptDiagnostics { - let tokenDeltas = diagnostics.compactMap(\.tokenDelta) + let tokenDeltas = contributions.compactMap(\.tokenDelta) let tokenTotals = tokenDeltas.isEmpty ? nil : tokenDiagnostics(tokenDeltas) - let contextSamples = diagnostics + let contextSamples = ( + contributions.compactMap { contribution in + contribution.contextUsage.map { + ( + contribution.date, + $0.totalTokens, + contribution.context?.modelContextWindow + ) + } + } + + diagnostics .filter { $0.key == .context } - .sorted { $0.date < $1.date } .compactMap { diagnostic -> (Date, Int64, Int64?)? in guard case let .tokens(usage) = diagnostic.value else { return nil @@ -899,6 +1393,8 @@ struct UsageReceiptSnapshot: Equatable, Sendable { diagnostic.context?.modelContextWindow ) } + ) + .sorted { $0.0 < $1.0 } let context = contextSamples.last.map { last in UsageReceiptContextDiagnostics( usedTokens: last.1, @@ -960,7 +1456,7 @@ struct UsageReceiptSnapshot: Equatable, Sendable { } else if tokenTotals?.reconciles == false { coverage = .low reason = "Token components do not match total" - } else if diagnostics.isEmpty { + } else if diagnostics.isEmpty && contributions.isEmpty { coverage = .unavailable reason = "No local diagnostics were observed" } else { @@ -974,7 +1470,10 @@ struct UsageReceiptSnapshot: Equatable, Sendable { tools: tools, compactions: compactions, duration: duration, - sources: Array(Set(diagnostics.map(\.source))) + sources: Array( + Set(diagnostics.map(\.source)) + .union(contributions.map(\.tokenSource)) + ) .sorted { $0.rawValue < $1.rawValue }, coverage: coverage, reason: reason @@ -990,7 +1489,7 @@ struct UsageReceiptSnapshot: Equatable, Sendable { _ value: KeyPath ) -> Int64? { guard values.allSatisfy({ - $0.observedComponents.contains(component) + $0.observes(component) }) else { return nil } @@ -1157,13 +1656,37 @@ struct UsageReceiptSnapshot: Equatable, Sendable { return (.notApplicable, "No Usage Receipts are available") } + private static func receiptEvidence( + _ receipts: [UsageReceiptSummary] + ) -> (CoverageLevel, String?) { + for level in [ + CoverageLevel.unavailable, + .low, + .partial, + .high, + .complete, + .notApplicable + ] { + if let receipt = receipts.first(where: { + $0.coverage == level + }) { + return (level, receipt.reason) + } + } + return (.notApplicable, "No Usage Receipts are available") + } + private func cumulativePoints( _ contributions: [Contribution] ) -> [LocalTokenActivityPoint] { var total: Int64 = 0 var points: [LocalTokenActivityPoint] = [] for contribution in contributions.sorted(by: { $0.date < $1.date }) { - total += contribution.tokens + let addition = total.addingReportingOverflow( + contribution.tokens + ) + guard !addition.overflow else { return [] } + total = addition.partialValue let point = LocalTokenActivityPoint( date: contribution.date, tokens: total @@ -1210,6 +1733,20 @@ enum UsageReceiptAggregator { taskIDsByRoot[rootTaskID, default: []].insert(taskID) } } + var sharedContexts: [ + LocalActivityContext: UsageReceiptSnapshot.SharedContext + ] = [:] + func sharedContext( + _ context: LocalActivityContext? + ) -> UsageReceiptSnapshot.SharedContext? { + guard let context else { return nil } + if let existing = sharedContexts[context] { + return existing + } + let shared = UsageReceiptSnapshot.SharedContext(context) + sharedContexts[context] = shared + return shared + } let timestampParser = LocalEventTimestampParser() var seen = Set() var contributions: [UsageReceiptSnapshot.Contribution] = [] @@ -1220,8 +1757,7 @@ enum UsageReceiptAggregator { seen.insert(eventID).inserted, let timestamp = fact.eventTimestamp, let date = timestampParser.date(from: timestamp), - date >= interval.start, - date < interval.end else { + date >= interval.start else { continue } let unboundedReasons = [ @@ -1258,16 +1794,17 @@ enum UsageReceiptAggregator { tokens: tokens, rootTaskID: rootID, projectLabel: projectLabel, - context: fact.context, + sharedContext: sharedContext(fact.context), tokenSource: fact.source.source, - hasUnboundedCounter: hasUnboundedCounter + hasUnboundedCounter: hasUnboundedCounter, + tokenDelta: fact.tokenDelta, + contextUsage: fact.contextUsage ) ) } var seenDiagnostics = Set() var diagnostics: [UsageReceiptSnapshot.DiagnosticContribution] = [] let diagnosticKeys: Set = [ - .token, .context, .time, .tool, @@ -1297,22 +1834,24 @@ enum UsageReceiptAggregator { let projectLabel = rootID.flatMap { projectionByTask[$0]?.projectLabel } + guard let payload = diagnosticPayload(for: fact) else { + continue + } let diagnostic = UsageReceiptSnapshot.DiagnosticContribution( date: date, rootTaskID: rootID, projectLabel: projectLabel, - context: fact.context, - key: fact.key, - value: fact.value, - tokenDelta: fact.tokenDelta, - availability: fact.availability, - reason: fact.reason, + sharedContext: sharedContext(fact.context), + payload: payload, eventID: eventID, source: fact.source.source ) - if diagnostic.intersects(interval) { - diagnostics.append(diagnostic) + guard diagnostic.intersects( + DateInterval(start: interval.start, end: .distantFuture) + ) else { + continue } + diagnostics.append(diagnostic) } return UsageReceiptSnapshot( contributions: contributions, @@ -1324,6 +1863,37 @@ enum UsageReceiptAggregator { ) } + private static func diagnosticPayload( + for fact: LocalActivityFact + ) -> UsageReceiptSnapshot.DiagnosticPayload? { + switch fact.key { + case .context: + guard case let .tokens(usage) = fact.value else { + return .context(nil) + } + return .context(usage) + case .time: + guard case let .turnTiming(timing) = fact.value else { + return .time(nil) + } + return .time(timing) + case .tool: + guard case let .text(toolClass) = fact.value else { + return .tool(nil) + } + return .tool(toolClass) + case .execution, .toolTime, .wait, .poll: + guard case let .duration(duration) = fact.value else { + return .duration(fact.key, nil) + } + return .duration(fact.key, duration) + case .compaction: + return .compaction + default: + return nil + } + } + private static func rootTaskID( for taskID: String, projections: [String: ThreadProjection], diff --git a/Tests/CodexLimitsTests/ActiveTimeAvailabilityTests.swift b/Tests/CodexLimitsTests/ActiveTimeAvailabilityTests.swift index d679ee6..4067d3c 100644 --- a/Tests/CodexLimitsTests/ActiveTimeAvailabilityTests.swift +++ b/Tests/CodexLimitsTests/ActiveTimeAvailabilityTests.swift @@ -413,18 +413,60 @@ final class ActiveTimeAvailabilityTests: XCTestCase { end: boundary.addingTimeInterval(1_800), eventDate: boundary.addingTimeInterval(1_800) ) - let index = LocalActivityFactIndex([missingEnd, missingStart]) + let facts = [missingEnd, missingStart] + let index = LocalActivityFactIndex(facts) XCTAssertEqual( - Set(index.activityFacts(in: first.interval).compactMap(\.eventID)), + Set( + index.activityFacts( + in: first.interval, + from: facts + ).compactMap(\.eventID) + ), ["missing-end", "missing-start"] ) XCTAssertEqual( - Set(index.activityFacts(in: second.interval).compactMap(\.eventID)), + Set( + index.activityFacts( + in: second.interval, + from: facts + ).compactMap(\.eventID) + ), ["missing-end", "missing-start"] ) } + func testActivityIndexDoesNotKeepTheFactBufferShared() { + let first = week(index: 0, movement: 80) + var facts = [ + timingFact( + id: "turn-1", + taskID: "task-1", + start: first.interval.start, + end: first.interval.start.addingTimeInterval(60) + ) + ] + facts.reserveCapacity(2) + let originalAddress = facts.withUnsafeBufferPointer(\.baseAddress) + let index = LocalActivityFactIndex(facts) + + facts.append( + timingFact( + id: "turn-2", + taskID: "task-2", + start: first.interval.start.addingTimeInterval(120), + end: first.interval.start.addingTimeInterval(180) + ) + ) + + withExtendedLifetime(index) { + XCTAssertEqual( + facts.withUnsafeBufferPointer(\.baseAddress), + originalAddress + ) + } + } + func testHistoricalSourceGapStaysLowOutsideItsObservationTime() { let affected = week( index: 0, diff --git a/Tests/CodexLimitsTests/ActivityTimelineTests.swift b/Tests/CodexLimitsTests/ActivityTimelineTests.swift index d7ec508..df7b588 100644 --- a/Tests/CodexLimitsTests/ActivityTimelineTests.swift +++ b/Tests/CodexLimitsTests/ActivityTimelineTests.swift @@ -589,6 +589,50 @@ final class ActivityTimelineTests: XCTestCase { reader.activeTimeAvailability.activeTimeThisWeek, 200 ) + XCTAssertEqual( + reader.activeTimeAvailability.maximumConcurrency, + 1 + ) + } + + func testAggregatorDoesNotRetainCompletedTurnsBeforeItsInterval() { + let interval = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 2_000) + ) + let snapshot = ActivityTimelineAggregator.evaluate( + facts: [ + timingFact( + eventID: "old", + taskID: "root", + turnID: "old", + start: 100, + end: 200 + ), + timingFact( + eventID: "current", + taskID: "root", + turnID: "current", + start: 1_100, + end: 1_200 + ) + ], + projections: [projection(taskID: "root", project: "atlas")], + interval: interval, + observation: .continuous( + sourceVersion: "0.145.0", + observedAt: interval.end + ) + ) + let wide = DateInterval( + start: Date(timeIntervalSince1970: 0), + end: interval.end + ) + + XCTAssertEqual( + snapshot.slice(in: wide, filters: .all).activeTime, + 100 + ) } private func accountSnapshot( diff --git a/Tests/CodexLimitsTests/CodexAssistedInsightTests.swift b/Tests/CodexLimitsTests/CodexAssistedInsightTests.swift index b602830..7bdb04c 100644 --- a/Tests/CodexLimitsTests/CodexAssistedInsightTests.swift +++ b/Tests/CodexLimitsTests/CodexAssistedInsightTests.swift @@ -554,6 +554,73 @@ final class CodexAssistedInsightTests: XCTestCase { XCTAssertTrue(savedResults.isEmpty) } + func testHistoryBeyondLegacySizeLimitRemainsReadableAndAppendable() async throws { + let root = temporaryDirectory() + let fileURL = root.appendingPathComponent("assisted.json") + let scope = analysisScope(accountPartitionID: "account-a") + let history = CodexAssistedHistory(fileURL: fileURL) + try await history.recordResult(analysisResult(), scope: scope) + + let handle = try FileHandle(forWritingTo: fileURL) + try handle.seekToEnd() + let whitespace = Data(repeating: 0x20, count: 1_024 * 1_024) + while try handle.offset() <= 64 * 1_024 * 1_024 { + try handle.write(contentsOf: whitespace) + } + try handle.close() + + let restarted = CodexAssistedHistory(fileURL: fileURL) + let restored = await restarted.results( + accountPartitionID: "account-a" + ) + try await restarted.recordResult( + analysisResult( + observedAt: Date(timeIntervalSince1970: 3_000) + ), + scope: scope + ) + + XCTAssertEqual(restored.count, 1) + let restoredAgain = await CodexAssistedHistory(fileURL: fileURL) + .results(accountPartitionID: "account-a") + XCTAssertEqual( + restoredAgain.map(\.result.observedAt), + [ + Date(timeIntervalSince1970: 2_012), + Date(timeIntervalSince1970: 3_000) + ] + ) + } + + func testCorruptDeletionMarkerFailsClosed() async throws { + let root = temporaryDirectory() + let fileURL = root.appendingPathComponent("assisted.json") + let markerURL = root.appendingPathComponent("deleting.json") + let scope = analysisScope(accountPartitionID: "account-a") + let first = CodexAssistedHistory( + fileURL: fileURL, + deletionMarkerURL: markerURL + ) + try await first.recordResult(analysisResult(), scope: scope) + let before = try Data(contentsOf: fileURL) + try Data("{not-json}".utf8).write(to: markerURL) + + let restarted = CodexAssistedHistory( + fileURL: fileURL, + deletionMarkerURL: markerURL + ) + let restored = await restarted.results( + accountPartitionID: "account-a" + ) + XCTAssertTrue(restored.isEmpty) + do { + try await restarted.recordResult(analysisResult(), scope: scope) + XCTFail("Expected the corrupt deletion marker to fail closed") + } catch {} + + XCTAssertEqual(try Data(contentsOf: fileURL), before) + } + func testDeletionMarkerSuppressesOldRecordsAndKeepsANewGeneration() async throws { let root = temporaryDirectory() let historyDirectory = root.appendingPathComponent( @@ -1018,6 +1085,31 @@ final class CodexAssistedInsightTests: XCTestCase { XCTAssertEqual(fixture.snapshot().connectionCount, 1) } + func testLiveClientRejectsOversizedCatalogLineAndUsesFreshConnection() async throws { + let fixture = CodexAssistedProtocolFixture( + oversizedCatalogLineOnFirstConnection: true, + maximumLineBytes: 1_024 + ) + let client = CodexAssistedClient( + makeConnection: { try fixture.makeConnection() }, + timeout: 1, + now: { fixture.now() } + ) + + do { + _ = try await client.eligibleProfile() + XCTFail("Expected the oversized line to close the connection") + } catch CodexAssistedClientError.connectionLost { + // Expected. + } catch { + XCTFail("Expected connectionLost, got \(error)") + } + + let selected = try await client.eligibleProfile() + XCTAssertNotNil(selected) + XCTAssertEqual(fixture.snapshot().connectionCount, 2) + } + func testLiveCatalogHidesTheActionWithoutASignedInAccount() async throws { let fixture = CodexAssistedProtocolFixture(accountIsMissing: true) let client = CodexAssistedClient( @@ -1441,6 +1533,8 @@ private final class CodexAssistedProtocolFixture: @unchecked Sendable { private let delaysThreadResponse: Bool private let accountIsMissing: Bool private let instructionSources: InstructionSourcesFixture + private let oversizedCatalogLineOnFirstConnection: Bool + private let maximumLineBytes: Int private var connections = 0 private var modelLists = 0 private var threadStarts = 0 @@ -1461,7 +1555,9 @@ private final class CodexAssistedProtocolFixture: @unchecked Sendable { addsUnknownEnabledFeature: Bool = false, delaysThreadResponse: Bool = false, accountIsMissing: Bool = false, - instructionSources: InstructionSourcesFixture = .empty + instructionSources: InstructionSourcesFixture = .empty, + oversizedCatalogLineOnFirstConnection: Bool = false, + maximumLineBytes: Int = 16 * 1_024 * 1_024 ) { self.sendsToolCall = sendsToolCall self.sendsUnknownItem = sendsUnknownItem @@ -1470,6 +1566,9 @@ private final class CodexAssistedProtocolFixture: @unchecked Sendable { self.delaysThreadResponse = delaysThreadResponse self.accountIsMissing = accountIsMissing self.instructionSources = instructionSources + self.oversizedCatalogLineOnFirstConnection = + oversizedCatalogLineOnFirstConnection + self.maximumLineBytes = maximumLineBytes } func now() -> Date { @@ -1500,7 +1599,10 @@ private final class CodexAssistedProtocolFixture: @unchecked Sendable { func makeConnection() throws -> CodexAppServerConnection { let requests = Pipe() let responses = Pipe() - lock.withLock { connections += 1 } + let connectionNumber = lock.withLock { + connections += 1 + return connections + } Task { for try await line in requests.fileHandleForReading.bytes.lines { guard let request = try? JSONSerialization.jsonObject( @@ -1524,7 +1626,12 @@ private final class CodexAssistedProtocolFixture: @unchecked Sendable { try? responses.fileHandleForWriting.close() return } - response = #"{"id":\#(id),"result":{"data":[{"id":"gpt-5.6-luna","model":"gpt-5.6-luna","displayName":"GPT-5.6 Luna","description":"","hidden":false,"isDefault":false,"defaultReasoningEffort":"medium","supportedReasoningEfforts":[{"reasoningEffort":"medium","description":"Balanced"}]}],"nextCursor":null}}"# + let padding = + oversizedCatalogLineOnFirstConnection + && connectionNumber == 1 + ? String(repeating: "x", count: maximumLineBytes) + : "" + response = #"{"id":\#(id),"result":{"data":[{"id":"gpt-5.6-luna","model":"gpt-5.6-luna","displayName":"GPT-5.6 Luna","description":"\#(padding)","hidden":false,"isDefault":false,"defaultReasoningEffort":"medium","supportedReasoningEfforts":[{"reasoningEffort":"medium","description":"Balanced"}]}],"nextCursor":null}}"# case "account/rateLimits/read": let read = lock.withLock { rateReads += 1 @@ -1624,6 +1731,7 @@ private final class CodexAssistedProtocolFixture: @unchecked Sendable { return CodexAppServerConnection( input: requests.fileHandleForWriting, output: responses.fileHandleForReading, + maximumLineBytes: maximumLineBytes, isRunning: { true }, stop: { try? requests.fileHandleForWriting.close() diff --git a/Tests/CodexLimitsTests/CodexClientTests.swift b/Tests/CodexLimitsTests/CodexClientTests.swift index da428a8..b86bd38 100644 --- a/Tests/CodexLimitsTests/CodexClientTests.swift +++ b/Tests/CodexLimitsTests/CodexClientTests.swift @@ -288,6 +288,47 @@ final class CodexClientTests: XCTestCase { XCTAssertEqual(server.threadListReadCount, 1) } + func testCancelledQueuedRequestDoesNotReachTheServer() async throws { + let server = PersistentAppServerFixture( + delaysFirstRateLimitResponse: true + ) + let client = CodexClient( + makeConnection: server.makeConnection, + timeout: 1 + ) + let fetch = Task { + try await client.fetch( + fetchedAt: Date(timeIntervalSince1970: 1_900_000) + ) + } + while server.rateLimitReadCount == 0 { + await Task.yield() + } + let queued = Task { + try await client.threadProjectionResponse( + for: .list( + cursor: nil, + limit: 100, + useStateDBOnly: true, + sortKey: "updated_at" + ) + ) + } + await Task.yield() + queued.cancel() + + _ = try await fetch.value + do { + _ = try await queued.value + XCTFail("Expected cancellation") + } catch is CancellationError { + // Expected. + } catch { + XCTFail("Expected CancellationError, got \(error)") + } + XCTAssertEqual(server.threadListReadCount, 0) + } + func testBrokenConnectionReconnectsWithoutLosingTheRefresh() async throws { let server = PersistentAppServerFixture(dropsFirstConnection: true) let client = CodexClient( @@ -304,6 +345,25 @@ final class CodexClientTests: XCTestCase { XCTAssertEqual(server.initializationCount, 2) } + func testOversizedProtocolLineReconnectsWithoutLosingTheRefresh() async throws { + let server = PersistentAppServerFixture( + oversizedInitializationOnFirstConnection: true, + maximumLineBytes: 1_024 + ) + let client = CodexClient( + makeConnection: server.makeConnection, + timeout: 1 + ) + + let result = try await client.fetch( + fetchedAt: Date(timeIntervalSince1970: 1_900_000) + ) + + XCTAssertEqual(result.snapshot.mainLimit?.window.remainingPercent, 80) + XCTAssertEqual(server.connectionCount, 2) + XCTAssertEqual(server.initializationCount, 2) + } + func testSparseRateLimitNotificationTriggersFullReconciliation() async throws { let server = PersistentAppServerFixture( sendsSparseRateLimitUpdate: true @@ -813,6 +873,35 @@ final class CodexClientTests: XCTestCase { ) } + func testInvalidSpendControlPercentFailsTheAccountRead() { + let usage = Data( + #"{"id":3,"result":{"dailyUsageBuckets":[]}}"#.utf8 + ) + for remaining in ["-1", "101", "1e308"] { + let rateLimits = Data( + #""" + {"id":2,"result":{"rateLimits":{ + "limitId":"codex", + "secondary":{"usedPercent":20,"windowDurationMins":10080,"resetsAt":2000000}, + "individualLimit":{"limit":"50","used":"10","remainingPercent":\#(remaining),"resetsAt":2100000} + }}} + """#.utf8 + ) + + XCTAssertThrowsError( + try CodexClient.decode( + rateLimitsResponse: rateLimits, + usageResponse: usage, + fetchedAt: Date(timeIntervalSince1970: 1_900_000) + ) + ) { error in + guard case CodexClientError.invalidResponse = error else { + return XCTFail("Expected invalidResponse, got \(error)") + } + } + } + } + func testNestedMissingFactsKeepTheirOwnObservationTimes() { let previousReadAt = Date(timeIntervalSince1970: 1_900_000) let currentReadAt = Date(timeIntervalSince1970: 1_900_060) @@ -948,6 +1037,45 @@ final class CodexClientTests: XCTestCase { } } + func testInvalidAllowanceWindowUsesTheIncompatibleResponseError() { + let usage = Data( + #"{"id":3,"result":{"dailyUsageBuckets":[]}}"#.utf8 + ) + for duration in [0, -10_080] { + let rateLimits = Data( + #"{"id":2,"result":{"rateLimits":{"limitId":"codex","secondary":{"usedPercent":20,"windowDurationMins":\#(duration),"resetsAt":2000000}}}}"# + .utf8 + ) + + XCTAssertThrowsError( + try CodexClient.decode( + rateLimitsResponse: rateLimits, + usageResponse: usage, + fetchedAt: Date(timeIntervalSince1970: 1_900_000) + ) + ) { error in + guard case CodexClientError.invalidResponse = error else { + return XCTFail("Expected invalidResponse, got \(error)") + } + } + } + } + + func testUsageWindowDecodingRejectsNonFiniteRemainingPercent() throws { + let decoder = JSONDecoder() + decoder.nonConformingFloatDecodingStrategy = .convertFromString( + positiveInfinity: "Infinity", + negativeInfinity: "-Infinity", + nan: "NaN" + ) + let data = Data( + #"{"remainingPercent":"NaN","resetsAt":2000000,"durationMinutes":10080}"# + .utf8 + ) + + XCTAssertThrowsError(try decoder.decode(UsageWindow.self, from: data)) + } + func testMalformedDailyTokenBucketsFailTheWholeAccountRead() { let rateLimits = Data(#""" {"id":2,"result":{ @@ -1096,6 +1224,8 @@ private final class PersistentAppServerFixture: @unchecked Sendable { private let includesChangingResetDetails: Bool private let errorBodiesByMethod: [String: String] private let initializationResultBody: String? + private let oversizedInitializationOnFirstConnection: Bool + private let maximumLineBytes: Int private var connections = 0 private var initializations = 0 private var rateLimitReads = 0 @@ -1146,7 +1276,9 @@ private final class PersistentAppServerFixture: @unchecked Sendable { initializeUserAgent: String? = nil, includesChangingResetDetails: Bool = false, errorBodiesByMethod: [String: String] = [:], - initializationResultBody: String? = nil + initializationResultBody: String? = nil, + oversizedInitializationOnFirstConnection: Bool = false, + maximumLineBytes: Int = 16 * 1_024 * 1_024 ) { self.dropsFirstConnection = dropsFirstConnection self.stallsFirstConnection = stallsFirstConnection @@ -1162,6 +1294,9 @@ private final class PersistentAppServerFixture: @unchecked Sendable { self.includesChangingResetDetails = includesChangingResetDetails self.errorBodiesByMethod = errorBodiesByMethod self.initializationResultBody = initializationResultBody + self.oversizedInitializationOnFirstConnection = + oversizedInitializationOnFirstConnection + self.maximumLineBytes = maximumLineBytes } func makeConnection() throws -> CodexAppServerConnection { @@ -1191,7 +1326,10 @@ private final class PersistentAppServerFixture: @unchecked Sendable { switch method { case "initialize": lock.withLock { initializations += 1 } - if let initializationResultBody { + if oversizedInitializationOnFirstConnection, + connectionNumber == 1 { + response = #"{"id":\#(id),"result":{"padding":"\#(String(repeating: "x", count: maximumLineBytes))"}}"# + } else if let initializationResultBody { response = #"{"id":\#(id),"result":\#(initializationResultBody)}"# } else if let initializeUserAgent { @@ -1318,6 +1456,7 @@ private final class PersistentAppServerFixture: @unchecked Sendable { return CodexAppServerConnection( input: requests.fileHandleForWriting, output: responses.fileHandleForReading, + maximumLineBytes: maximumLineBytes, isRunning: { true }, stop: { try? requests.fileHandleForWriting.close() diff --git a/Tests/CodexLimitsTests/ForecastEngineTests.swift b/Tests/CodexLimitsTests/ForecastEngineTests.swift index a5a7e44..3b4147f 100644 --- a/Tests/CodexLimitsTests/ForecastEngineTests.swift +++ b/Tests/CodexLimitsTests/ForecastEngineTests.swift @@ -91,6 +91,69 @@ final class ForecastEngineTests: XCTestCase { XCTAssertEqual(result.historicalReferenceSource, .tokenEstimate) } + func testLargeTokenBucketsDoNotOverflowTheForecast() { + let day: TimeInterval = 86_400 + let now = Date(timeIntervalSince1970: 100 * day) + let reset = now.addingTimeInterval(2 * day) + let currentDate = now.addingTimeInterval(-5 * day) + let window = UsageWindow( + remainingPercent: 90, + resetsAt: reset, + durationMinutes: 7 * 24 * 60 + ) + let tokenHistory = [ + TokenDay(date: currentDate, tokens: Int64.max), + TokenDay(date: currentDate, tokens: Int64.max), + TokenDay( + date: now.addingTimeInterval(-10 * day), + tokens: 1 + ) + ] + + let result = ForecastEngine.evaluate( + window: window, + samples: [], + tokenHistory: tokenHistory, + safetyBuffer: 3, + now: now, + previousStatus: nil + ) + + XCTAssertTrue(result.currentPercentPerDay.isFinite) + XCTAssertTrue(result.expectedRemainingAtReset.isFinite) + XCTAssertEqual(result.historicalReferenceSource, .tokenEstimate) + } + + func testOutOfRangeTokenDateDoesNotTrapTheForecast() { + let day: TimeInterval = 86_400 + let now = Date(timeIntervalSince1970: 100 * day) + let window = UsageWindow( + remainingPercent: 90, + resetsAt: now.addingTimeInterval(2 * day), + durationMinutes: 7 * 24 * 60 + ) + + let result = ForecastEngine.evaluate( + window: window, + samples: [], + tokenHistory: [ + TokenDay( + date: Date( + timeIntervalSince1970: Double.greatestFiniteMagnitude + ), + tokens: 1 + ) + ], + safetyBuffer: 3, + now: now, + previousStatus: nil + ) + + XCTAssertTrue(result.currentPercentPerDay.isFinite) + XCTAssertTrue(result.expectedRemainingAtReset.isFinite) + XCTAssertNil(result.historicalReferenceSource) + } + func testOnlyCompleteHighCoverageWeeklyWindowsBecomeAccountHistory() { let day: TimeInterval = 86_400 let halfHour: TimeInterval = 30 * 60 diff --git a/Tests/CodexLimitsTests/LocalActivityCollectorTests.swift b/Tests/CodexLimitsTests/LocalActivityCollectorTests.swift index 8f16be0..6591098 100644 --- a/Tests/CodexLimitsTests/LocalActivityCollectorTests.swift +++ b/Tests/CodexLimitsTests/LocalActivityCollectorTests.swift @@ -126,9 +126,13 @@ final class LocalActivityCollectorTests: XCTestCase { interval: interval, observedAt: observedAt ) + let failedProjectionSource = ReadOnlyThreadProjectionSource { _ in + throw CocoaError(.fileReadUnknown) + } let restarted = LocalActivityCollector( rootDirectory: fixture.root, - stateDirectory: stateDirectory + stateDirectory: stateDirectory, + projectionSource: failedProjectionSource ) await restarted.selectPartition("stable-account") let afterRestart = await restarted.refresh( @@ -153,6 +157,52 @@ final class LocalActivityCollectorTests: XCTestCase { ) } + func testRewrittenRolloutKeepsFactsFromUnchangedActiveFiles() async throws { + let fixture = try CollectorFixture() + let rewritten = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-2", + lines: [ + fixture.session(threadID: "task-2", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 3), + fixture.tokens(total: 1_000, ordinal: 2, minute: 4) + ] + ) + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: fixture.root.appendingPathComponent("state") + ) + await collector.selectPartition("stable-account") + let interval = try fixture.interval() + _ = await collector.refresh(interval: interval) + try Data( + ( + [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 800, ordinal: 2, minute: 2) + ].joined(separator: "\n") + "\n" + ).utf8 + ).write(to: rewritten, options: .atomic) + + let result = await collector.refresh(interval: interval) + + XCTAssertEqual( + result.facts.filter { $0.key == .token } + .compactMap(\.numericDelta).sorted(), + [700, 900] + ) + } + func testNewestRolloutDiscontinuityReplacesAnOlderBoundary() async throws { let fixture = try CollectorFixture() let file = try fixture.rollout( @@ -365,6 +415,121 @@ final class LocalActivityCollectorTests: XCTestCase { ) } + func testRestoreCompactsLegacyContextFactsIntoTheirTokenFacts() async throws { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + #"{"timestamp":"2026-07-28T10:01:00.000Z","ordinal":1,"type":"event_msg","payload":{"type":"token_count","model_context_window":272000,"info":{"total_token_usage":{"total_tokens":100},"last_token_usage":{"total_tokens":80}}}}"#, + #"{"timestamp":"2026-07-28T10:02:00.000Z","ordinal":2,"type":"event_msg","payload":{"type":"token_count","model_context_window":272000,"info":{"total_token_usage":{"total_tokens":600},"last_token_usage":{"total_tokens":90}}}}"# + ] + ) + let stateDirectory = fixture.root.appendingPathComponent( + "collector-state", + isDirectory: true + ) + let first = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await first.selectPartition("stable-account") + _ = await first.refresh(interval: try fixture.interval()) + let partitionDirectory = stateDirectory.appendingPathComponent( + "stable-account", + isDirectory: true + ) + let factsFile = try XCTUnwrap( + FileManager.default.contentsOfDirectory( + at: partitionDirectory, + includingPropertiesForKeys: nil + ).first { $0.lastPathComponent.hasSuffix(".facts.jsonl") } + ) + let decoder = JSONDecoder() + let encoder = JSONEncoder() + let persisted = try String(contentsOf: factsFile, encoding: .utf8) + .split(separator: "\n") + .map { try decoder.decode(LocalActivityFact.self, from: Data($0.utf8)) } + var legacy: [LocalActivityFact] = [] + for fact in persisted { + guard fact.key == .token, + let contextUsage = fact.contextUsage else { + legacy.append(fact) + continue + } + legacy.append( + LocalActivityFact( + key: fact.key, + availability: fact.availability, + value: fact.value, + numericDelta: fact.numericDelta, + tokenSegment: fact.tokenSegment, + reason: fact.reason, + eventID: fact.eventID, + eventTimestamp: fact.eventTimestamp, + source: fact.source, + context: fact.context, + tokenDelta: fact.tokenDelta + ) + ) + legacy.append( + LocalActivityFact( + key: .context, + availability: .available, + value: .tokens(contextUsage), + numericDelta: nil, + tokenSegment: nil, + reason: nil, + eventID: fact.eventID, + eventTimestamp: fact.eventTimestamp, + source: contextUsage.totalTokens == 90 + ? LocalActivitySourceMetadata( + source: fact.source.source, + sourceVersion: fact.source.sourceVersion, + schemaVersion: fact.source.schemaVersion, + sourceGeneration: + fact.source.sourceGeneration + 1, + historyMode: fact.source.historyMode, + observedAt: fact.source.observedAt + ) + : fact.source, + context: fact.context + ) + ) + } + var encoded = Data() + for fact in legacy { + encoded.append(try encoder.encode(fact)) + encoded.append(0x0A) + } + try encoded.write(to: factsFile) + + let restarted = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await restarted.selectPartition("stable-account") + let restored = await restarted.refresh(interval: try fixture.interval()) + + XCTAssertEqual( + restored.facts.filter { $0.key == .token } + .compactMap(\.contextUsage).map(\.totalTokens), + [80] + ) + XCTAssertEqual( + restored.facts.filter { + $0.key == .context && $0.availability == .available + }.compactMap { fact -> Int64? in + guard case let .tokens(usage) = fact.value else { + return nil + } + return usage.totalTokens + }, + [90] + ) + } + func testRestartedCollectorKeepsCursorAfterRolloutRename() async throws { let fixture = try CollectorFixture() let rollout = try fixture.rollout( @@ -934,7 +1099,7 @@ final class LocalActivityCollectorTests: XCTestCase { ) XCTAssertEqual(receiptSlice.receipts.first?.rootTaskID, "task-1") XCTAssertEqual(receiptSlice.receipts.first?.tokens, 500) - XCTAssertEqual(migrated["version"] as? Int, 6) + XCTAssertEqual(migrated["version"] as? Int, 7) XCTAssertNil(migrated["path"]) XCTAssertNotNil(migrated["pathFingerprint"]) } @@ -1167,7 +1332,7 @@ final class LocalActivityCollectorTests: XCTestCase { ) await collector.selectPartition("stable-account") let interval = try fixture.interval() - _ = await collector.refresh(interval: interval) + let first = await collector.refresh(interval: interval) let partitionDirectory = stateDirectory.appendingPathComponent( "stable-account", isDirectory: true @@ -1184,12 +1349,18 @@ final class LocalActivityCollectorTests: XCTestCase { fixture.tokens(total: 800, ordinal: 3, minute: 3), to: rollout ) - _ = await collector.refresh(interval: interval) + let appended = await collector.refresh(interval: interval) let after = try Data(contentsOf: factsFile) XCTAssertTrue(after.starts(with: before)) XCTAssertGreaterThan(after.count, before.count) XCTAssertLessThan(after.count - before.count, before.count) + XCTAssertNotEqual(appended.contentRevision, first.contentRevision) + XCTAssertEqual( + appended.facts.filter { $0.key == .token } + .compactMap(\.numericDelta), + [500, 200] + ) } func testIdleRefreshDoesNotRewriteDurableState() async throws { @@ -1242,6 +1413,173 @@ final class LocalActivityCollectorTests: XCTestCase { XCTAssertEqual(after, before) } + func testIdleRefreshKeepsContentRevision() async throws { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: fixture.root.appendingPathComponent("state") + ) + await collector.selectPartition("stable-account") + let interval = try fixture.interval() + + let first = await collector.refresh( + interval: interval, + observedAt: Date(timeIntervalSince1970: 100) + ) + let idle = await collector.refresh( + interval: interval, + observedAt: Date(timeIntervalSince1970: 200) + ) + + XCTAssertEqual(idle.contentRevision, first.contentRevision) + XCTAssertEqual(idle.facts, first.facts) + } + + func testProjectionObservationTimeDoesNotChangeContentRevision() async throws { + let fixture = try CollectorFixture() + let file = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let source = ReadOnlyThreadProjectionSource { request in + guard case .list(cursor: nil, _, _, _) = request else { + throw CocoaError(.fileReadUnknown) + } + return Data(#""" + {"result":{"data":[{ + "id":"task-1", + "parentThreadId":null, + "cliVersion":"0.145.0", + "cwd":"/synthetic/projects/atlas", + "path":"\#(file.path)", + "createdAt":1785232800, + "updatedAt":1785232920 + }],"nextCursor":null}} + """#.utf8) + } + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: fixture.root.appendingPathComponent("state"), + projectionSource: source + ) + await collector.selectPartition("stable-account") + let interval = try fixture.interval() + let first = await collector.refresh(interval: interval) + try await Task.sleep(nanoseconds: 1_000_000) + + let idle = await collector.refresh(interval: interval) + + XCTAssertNotEqual( + first.projections.first?.source.observedAt, + idle.projections.first?.source.observedAt + ) + XCTAssertEqual(idle.contentRevision, first.contentRevision) + } + + func testMissingRootKeepsLastPublishedFacts() async throws { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: fixture.root.appendingPathComponent("state") + ) + await collector.selectPartition("stable-account") + let interval = try fixture.interval() + let first = await collector.refresh(interval: interval) + + try FileManager.default.removeItem(at: fixture.root) + let unavailable = await collector.refresh(interval: interval) + + XCTAssertEqual(unavailable.facts, first.facts) + XCTAssertEqual(unavailable.contentRevision, 0) + XCTAssertEqual(unavailable.observation.coverage, .unavailable) + } + + func testIdleRefreshWithLaterIntervalEndReusesPublishedFacts() async throws { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: fixture.root.appendingPathComponent("state") + ) + await collector.selectPartition("stable-account") + let firstInterval = try fixture.interval( + end: "2026-07-28T12:00:00Z" + ) + let laterInterval = try fixture.interval( + end: "2026-07-28T13:00:00Z" + ) + + let first = await collector.refresh(interval: firstInterval) + let idle = await collector.refresh(interval: laterInterval) + + XCTAssertEqual(idle.bytesRead, 0) + XCTAssertEqual(idle.contentRevision, first.contentRevision) + XCTAssertEqual(idle.facts, first.facts) + } + + func testChangingIntervalStartDoesNotReusePriorIntervalFacts() async throws { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: fixture.root.appendingPathComponent("state") + ) + await collector.selectPartition("stable-account") + + _ = await collector.refresh(interval: try fixture.interval()) + let later = await collector.refresh( + interval: try fixture.interval( + start: "2026-07-28T11:00:00Z" + ) + ) + + XCTAssertTrue( + later.facts.filter { + $0.key == .token && $0.availability == .available + }.isEmpty + ) + XCTAssertEqual(later.bytesRead, 0) + } + func testRestartDoesNotLoadPersistedFactsOutsideTheCurrentInterval() async throws { let fixture = try CollectorFixture() _ = try fixture.rollout( @@ -1281,6 +1619,95 @@ final class LocalActivityCollectorTests: XCTestCase { XCTAssertEqual(collection.observation.coverage, .high) } + func testOversizedMetadataIsIgnoredAndRolloutIsRebuilt() async throws { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let stateDirectory = fixture.root.appendingPathComponent( + "collector-state", + isDirectory: true + ) + let first = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await first.selectPartition("stable-account") + let interval = try fixture.interval() + _ = await first.refresh(interval: interval) + let partitionDirectory = stateDirectory.appendingPathComponent( + "stable-account", + isDirectory: true + ) + let metadata = try XCTUnwrap( + FileManager.default.contentsOfDirectory( + at: partitionDirectory, + includingPropertiesForKeys: nil + ).first { $0.pathExtension == "json" } + ) + try Data(repeating: 0, count: 1_048_577).write( + to: metadata, + options: .atomic + ) + + let restarted = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await restarted.selectPartition("stable-account") + let rebuilt = await restarted.refresh(interval: interval) + + XCTAssertGreaterThan(rebuilt.bytesRead, 0) + XCTAssertEqual( + rebuilt.facts.filter { $0.key == .token } + .compactMap(\.numericDelta), + [500] + ) + XCTAssertEqual( + rebuilt.observation.reason, + "Saved local activity could not be read" + ) + } + + func testChangingIntervalsReloadsPersistedFactsWithoutRescanningRollout() async throws { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "past-task", + lines: [ + fixture.session(threadID: "past-task", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: fixture.root.appendingPathComponent("state") + ) + await collector.selectPartition("stable-account") + let past = try fixture.interval() + let first = await collector.refresh(interval: past) + let current = try fixture.interval( + start: "2026-08-04T00:00:00Z", + end: "2026-08-05T00:00:00Z" + ) + + let empty = await collector.refresh(interval: current) + let restored = await collector.refresh(interval: past) + + XCTAssertTrue(empty.facts.isEmpty) + XCTAssertEqual(restored.facts, first.facts) + XCTAssertEqual(restored.bytesRead, 0) + XCTAssertNotEqual(empty.contentRevision, first.contentRevision) + XCTAssertNotEqual(restored.contentRevision, empty.contentRevision) + } + func testCorruptPersistedFactsRebuildFromTheRollout() async throws { let fixture = try CollectorFixture() _ = try fixture.rollout( @@ -1330,6 +1757,62 @@ final class LocalActivityCollectorTests: XCTestCase { XCTAssertEqual(rebuilt.observation.coverage, .high) } + func testOversizedPersistedFactRebuildsFromTheRollout() async throws { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let stateDirectory = fixture.root.appendingPathComponent("state") + let first = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await first.selectPartition("stable-account") + let interval = try fixture.interval() + _ = await first.refresh(interval: interval) + let partition = stateDirectory.appendingPathComponent("stable-account") + let factsFile = try XCTUnwrap( + FileManager.default.contentsOfDirectory( + at: partition, + includingPropertiesForKeys: nil + ).first { $0.lastPathComponent.hasSuffix(".facts.jsonl") } + ) + let firstLine = try XCTUnwrap( + String(contentsOf: factsFile, encoding: .utf8) + .split(separator: "\n").first + ) + let padding = String( + repeating: "x", + count: BoundedJSONLReader.maximumRecordBytes + ) + let oversized = firstLine.replacingOccurrences( + of: "{", + with: #"{"padding":"\#(padding)","#, + options: [.anchored] + ) + "\n" + try Data(oversized.utf8).write(to: factsFile) + + let restarted = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await restarted.selectPartition("stable-account") + let restored = await restarted.refresh(interval: interval) + + XCTAssertEqual( + restored.facts.filter { $0.key == .token } + .compactMap(\.numericDelta), + [500] + ) + XCTAssertGreaterThan(restored.bytesRead, 0) + } + func testDeletedHistoryDoesNotReturnOnTheNextRefreshOrRestart() async throws { let fixture = try CollectorFixture() let rollout = try fixture.rollout( @@ -1603,6 +2086,43 @@ final class LocalActivityCollectorTests: XCTestCase { ) } + func testCancelledProjectionReadDoesNotScanOrSaveRollouts() async throws { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let stateDirectory = fixture.root.appendingPathComponent("state") + let delay = CancellableProjectionDelay() + let source = ReadOnlyThreadProjectionSource { _ in + try await delay.response() + } + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory, + projectionSource: source + ) + await collector.selectPartition("stable-account") + let refresh = Task { + await collector.refresh(interval: try fixture.interval()) + } + await delay.waitUntilStarted() + + refresh.cancel() + let result = try await refresh.value + + XCTAssertEqual(result.observation.coverage, .unavailable) + XCTAssertTrue(result.facts.isEmpty) + XCTAssertFalse( + FileManager.default.fileExists(atPath: stateDirectory.path) + ) + } + func testDurableWriteFailureLowersCoverage() async throws { let fixture = try CollectorFixture() _ = try fixture.rollout( @@ -1636,6 +2156,153 @@ final class LocalActivityCollectorTests: XCTestCase { [500] ) } + + func testResumableRestoreKeepsFactsFromEveryCandidateFile() async throws { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "older-task", + lines: [ + fixture.session(threadID: "older-task", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + var newerLines = [ + fixture.session(threadID: "newer-task", ordinal: 0) + ] + newerLines.reserveCapacity(10_101) + for ordinal in 1 ... 10_100 { + newerLines.append( + fixture.tokens(total: ordinal * 100, ordinal: ordinal, minute: 1) + ) + } + _ = try fixture.rollout( + day: "2026/07/29", + threadID: "newer-task", + lines: newerLines + ) + let interval = try fixture.interval( + end: "2026-07-30T00:00:00Z" + ) + let stateDirectory = fixture.root.appendingPathComponent("state") + let first = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await first.selectPartition("stable-account") + for _ in 0 ..< 10 { + _ = await first.refresh(interval: interval) + if await first.hasPendingImport() == false { break } + } + let firstPending = await first.hasPendingImport() + XCTAssertFalse(firstPending) + + let failedProjectionSource = ReadOnlyThreadProjectionSource { _ in + throw CocoaError(.fileReadUnknown) + } + let restarted = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory, + projectionSource: failedProjectionSource + ) + await restarted.selectPartition("stable-account") + let narrowInterval = try fixture.interval( + start: "2026-07-28T11:00:00Z", + end: "2026-07-30T00:00:00Z" + ) + var restored = await restarted.refresh(interval: narrowInterval) + let firstRestoreRevision = restored.contentRevision + let pendingAfterFirstRestore = await restarted.hasPendingImport() + XCTAssertTrue(pendingAfterFirstRestore) + for _ in 0 ..< 10 { + guard await restarted.hasPendingImport() else { break } + restored = await restarted.refresh( + interval: interval, + refreshMetadata: false + ) + } + + let pendingAfterRestore = await restarted.hasPendingImport() + XCTAssertFalse(pendingAfterRestore) + XCTAssertNotEqual(restored.contentRevision, firstRestoreRevision) + XCTAssertEqual( + restored.observation.reason, + "Local task discovery is incomplete" + ) + XCTAssertEqual( + restored.facts + .filter { $0.key == .token } + .compactMap(\.numericDelta) + .reduce(0, +), + 1_010_400 + ) + } + + func testNewerRefreshSupersedesASuspendedRefresh() async throws { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let delay = SupersededProjectionDelay() + let source = ReadOnlyThreadProjectionSource { request in + await delay.response(to: request) + } + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: fixture.root.appendingPathComponent("state"), + projectionSource: source + ) + let interval = try fixture.interval() + let stale = Task { + await collector.refresh(interval: interval) + } + await delay.waitUntilListStarted() + + let latest = await collector.refresh( + interval: interval, + refreshMetadata: false + ) + await delay.releaseList() + let superseded = await stale.value + let idle = await collector.refresh( + interval: interval, + refreshMetadata: false + ) + + XCTAssertEqual(latest.facts, idle.facts) + XCTAssertEqual( + latest.facts.filter { $0.key == .token }.compactMap(\.numericDelta), + [500] + ) + XCTAssertEqual(superseded.observation.coverage, .unavailable) + let pending = await collector.hasPendingImport() + XCTAssertFalse(pending) + } + + func testLowerCoverageDoesNotChangeTheFactCacheRevision() { + let collection = LocalActivityCollection( + facts: [], + projections: [], + observation: .continuous( + sourceVersion: "0.145.0", + observedAt: Date(timeIntervalSince1970: 100) + ), + bytesRead: 0, + contentRevision: 7 + ) + + XCTAssertEqual( + collection.loweringCoverage("Identity unavailable").contentRevision, + 7 + ) + } } private actor CollectorProjectionRequests { @@ -1719,6 +2386,61 @@ private actor ProjectionDelay { } } +private actor CancellableProjectionDelay { + private var started = false + + func response() async throws -> Data { + started = true + try await Task.sleep(nanoseconds: 60_000_000_000) + return Data() + } + + func waitUntilStarted() async { + while !started { + await Task.yield() + } + } +} + +private actor SupersededProjectionDelay { + private var listStarted = false + private var listContinuation: CheckedContinuation? + + func response(to request: ThreadProjectionReadRequest) async -> Data { + switch request { + case .list: + listStarted = true + return await withCheckedContinuation { continuation in + listContinuation = continuation + } + case let .read(threadID, _): + return Data(#""" + {"result":{"thread":{ + "id":"\#(threadID)", + "parentThreadId":null, + "cliVersion":"0.145.0", + "cwd":"/synthetic/project", + "createdAt":1785232800, + "updatedAt":1785232920 + }}} + """#.utf8) + } + } + + func waitUntilListStarted() async { + while !listStarted { + await Task.yield() + } + } + + func releaseList() { + listContinuation?.resume( + returning: Data(#"{"result":{"data":[],"nextCursor":null}}"#.utf8) + ) + listContinuation = nil + } +} + private final class CollectorFixture { let root: URL diff --git a/Tests/CodexLimitsTests/LocalActivityNormalizerTests.swift b/Tests/CodexLimitsTests/LocalActivityNormalizerTests.swift index 57e7a81..e7ca512 100644 --- a/Tests/CodexLimitsTests/LocalActivityNormalizerTests.swift +++ b/Tests/CodexLimitsTests/LocalActivityNormalizerTests.swift @@ -75,22 +75,53 @@ final class LocalActivityNormalizerTests: XCTestCase { ) ) XCTAssertEqual( - second.facts(.context).last?.value, - .tokens( - LocalTokenUsage( - inputTokens: 800, - cachedInputTokens: 300, - cacheWriteInputTokens: 15, - outputTokens: 140, - reasoningOutputTokens: 60, - totalTokens: 940 - ) + second.facts(.token).last?.contextUsage, + LocalTokenUsage( + inputTokens: 800, + cachedInputTokens: 300, + cacheWriteInputTokens: 15, + outputTokens: 140, + reasoningOutputTokens: 60, + totalTokens: 940 ) ) XCTAssertEqual( - second.facts(.context).last?.context?.modelContextWindow, + second.facts(.token).last?.context?.modelContextWindow, 272_000 ) + XCTAssertFalse( + second.facts(.context).contains { + $0.availability == .available + } + ) + } + + func testContextSampleWithoutTokenCounterRemainsStandalone() { + let usage = LocalTokenUsage( + inputTokens: 800, + cachedInputTokens: 300, + cacheWriteInputTokens: 15, + outputTokens: 140, + reasoningOutputTokens: 60, + totalTokens: 940 + ) + let result = LocalActivityNormalizer().normalize( + records: [ + record( + id: "context-only", + type: "event_msg", + threadID: "task-root", + contextTokenUsage: usage + ) + ], + sourceGeneration: 0, + observedAt: Date(timeIntervalSince1970: 1_200) + ) + + XCTAssertEqual( + result.facts(.context).last?.value, + .tokens(usage) + ) } func testCompactionKeepsTurnContext() { @@ -335,6 +366,93 @@ final class LocalActivityNormalizerTests: XCTestCase { XCTAssertEqual(result.fact(.time)?.reason, "turn-start-not-observed") } + func testExtremeRestoredTokenStateDoesNotOverflow() { + let previous = LocalActivityNormalizationState( + sourceGeneration: 0, + sourceVersion: "0.145.0", + historyMode: "paginated", + lastTotalTokens: .min, + tokenSegment: .max + ) + + let result = LocalActivityNormalizer().normalize( + records: [ + record( + id: "extreme", + type: "event_msg", + threadID: "task-root", + tokens: .max + ) + ], + sourceGeneration: 0, + observedAt: Date(timeIntervalSince1970: 1_200), + previousState: previous + ) + + XCTAssertNil(result.facts(.token).last?.numericDelta) + XCTAssertEqual( + result.facts(.token).last?.reason, + "cumulative-counter-decreased" + ) + XCTAssertEqual(result.state.tokenSegment, .max) + } + + func testChunkedNormalizationMatchesOnePassFacts() { + let observedAt = Date(timeIntervalSince1970: 1_200) + let records = [ + record( + id: "session", + type: "session_meta", + threadID: "task-root" + ), + record( + id: "baseline", + type: "event_msg", + threadID: "task-root", + turnID: "turn-1", + model: "gpt-5.6-sol", + reasoning: "high", + tokens: 100 + ), + record( + id: "delta", + type: "event_msg", + threadID: "task-root", + tokens: 600 + ) + ] + let normalizer = LocalActivityNormalizer() + let onePass = normalizer.normalize( + records: records, + sourceGeneration: 0, + observedAt: observedAt + ) + let first = normalizer.normalize( + records: Array(records.prefix(2)), + sourceGeneration: 0, + observedAt: observedAt + ) + let second = normalizer.normalize( + records: Array(records.dropFirst(2)), + sourceGeneration: 0, + observedAt: observedAt, + previousState: first.state + ) + let chunkedFacts = (first.facts + second.facts).filter { + $0.eventID != nil + } + + XCTAssertEqual( + chunkedFacts, + onePass.facts.filter { $0.eventID != nil } + ) + XCTAssertEqual( + chunkedFacts.filter { $0.key == .token } + .compactMap(\.numericDelta), + [500] + ) + } + private func record( id: String, type: String, diff --git a/Tests/CodexLimitsTests/LocalActivityPerformanceTests.swift b/Tests/CodexLimitsTests/LocalActivityPerformanceTests.swift index 4ea5e0c..2803429 100644 --- a/Tests/CodexLimitsTests/LocalActivityPerformanceTests.swift +++ b/Tests/CodexLimitsTests/LocalActivityPerformanceTests.swift @@ -4,6 +4,146 @@ import XCTest @testable import CodexLimits final class LocalActivityPerformanceTests: XCTestCase { + func testStableUsageEvaluationReusesLargeLocalHistory() throws { + let now = Date(timeIntervalSince1970: 2_000_000) + let window = UsageWindow( + remainingPercent: 80, + resetsAt: now.addingTimeInterval(2 * 86_400), + durationMinutes: UsageHistoryPolicy.weeklyDurationMinutes + ) + let account = UsageSnapshot( + mainLimit: LimitReading( + limitId: "codex", + name: "Codex", + window: window + ), + otherLimits: [], + tokenHistory: [], + emergencyResetCount: 0, + fetchedAt: now + ) + let samples = [ + UsageSample( + observedAt: window.startsAt, + remainingPercent: 100, + resetsAt: window.resetsAt, + lifetimeTokens: 0 + ), + UsageSample( + observedAt: now, + remainingPercent: 80, + resetsAt: window.resetsAt, + lifetimeTokens: 100_000 + ) + ] + let source = LocalActivitySourceMetadata( + source: .rolloutJSONL, + sourceVersion: "0.145.0", + schemaVersion: "rollout-jsonl-v1", + sourceGeneration: 0, + historyMode: nil, + observedAt: now + ) + let timestamp = ISO8601DateFormatter().string( + from: now.addingTimeInterval(-60) + ) + let facts = (0 ..< 100_000).map { + LocalActivityFact( + key: .token, + availability: .available, + value: nil, + numericDelta: 1, + tokenSegment: 0, + reason: nil, + eventID: "token-\($0)", + eventTimestamp: timestamp, + source: source + ) + } + let observation = LocalActivityObservation.continuous( + sourceVersion: "0.145.0", + observedAt: now + ) + let compatibleSources: Set = [ + LocalTokenDefinitionSource(source) + ] + let first = UsageIntelligenceEngine.evaluate( + UsageIntelligenceInput( + account: account, + samples: samples, + safetyBuffer: 3, + sourceState: .available, + now: now, + previousStatus: nil, + accountPartitionID: "account-a", + localActivityFacts: [], + localActivityHistoryFacts: facts, + localActivityObservation: observation, + localActivityContentRevision: 7, + compatibleTokenSources: compatibleSources + ) + ) + let start = ProcessInfo.processInfo.systemUptime + + let refreshed = UsageIntelligenceEngine.evaluate( + UsageIntelligenceInput( + account: account, + samples: samples, + safetyBuffer: 3, + sourceState: .available, + now: now, + previousStatus: nil, + accountPartitionID: "account-a", + localActivityFacts: [], + localActivityHistoryFacts: facts, + localActivityObservation: observation, + localActivityContentRevision: 7, + reusableLocalAggregates: + try XCTUnwrap(first.reusableLocalAggregates), + compatibleTokenSources: compatibleSources + ) + ) + let milliseconds = + (ProcessInfo.processInfo.systemUptime - start) * 1_000 + + print( + String( + format: "STABLE_USAGE_EVALUATION facts=%d elapsed_ms=%.3f", + facts.count, + milliseconds + ) + ) + XCTAssertEqual( + refreshed.usagePerToken.current?.localTokenActivity, + 100_000 + ) + XCTAssertLessThan(milliseconds, 100) + } + + func testTimestampParsingStaysResponsive() { + let parser = LocalEventTimestampParser() + let timestamp = "2026-07-29T20:46:09.123Z" + let iterations = 50_000 + let start = ProcessInfo.processInfo.systemUptime + var parsed: Date? + + for _ in 0..= residentBefore + ? residentAfter - residentBefore + : 0 print( - String( - format: "PERSISTED_FACT_RESTORE records=%d wall_ms=%.3f", - restored.facts.count, - milliseconds - ) + [ + "PERSISTED_FACT_RESTORE", + "records=\(restored.facts.count)", + "refreshes=\(refreshCount)", + String( + format: "max_refresh_ms=%.3f", + maximumRefreshMilliseconds + ), + "resident_delta_bytes=\(residentDelta)" + ].joined(separator: " ") ) - XCTAssertEqual(restored.bytesRead, 0) + XCTAssertLessThanOrEqual(restored.bytesRead, 4_096) XCTAssertEqual( restored.facts.filter { $0.key == .token } .compactMap(\.numericDelta).count, recordCount - 1 ) - XCTAssertLessThan(milliseconds, 3_000) + XCTAssertEqual( + restored.facts.filter { $0.key == .token } + .compactMap(\.contextUsage).count, + recordCount + ) + XCTAssertFalse( + restored.facts.contains { + $0.key == .context && $0.availability == .available + } + ) + XCTAssertGreaterThan(refreshCount, 1) + let pendingAfterRestore = await restarted.hasPendingImport() + XCTAssertFalse(pendingAfterRestore) + XCTAssertLessThan(maximumRefreshMilliseconds, 3_000) + XCTAssertLessThan(residentDelta, 256 * 1_024 * 1_024) } func testRepresentativeFixtureMetrics() throws { diff --git a/Tests/CodexLimitsTests/LocalTokenActivityTests.swift b/Tests/CodexLimitsTests/LocalTokenActivityTests.swift index cc867af..f9b9009 100644 --- a/Tests/CodexLimitsTests/LocalTokenActivityTests.swift +++ b/Tests/CodexLimitsTests/LocalTokenActivityTests.swift @@ -76,6 +76,37 @@ final class LocalTokenActivityTests: XCTestCase { XCTAssertEqual(activity.reason, "Local task records are missing") } + func testCachedActivityUsesTheLatestCoverageObservation() { + let interval = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 2_000) + ) + let cached = LocalTokenActivityAggregator.evaluate( + facts: [tokenFact(id: "first", timestamp: 1_100, delta: 100)], + interval: interval, + observation: .continuous( + sourceVersion: "0.145.0", + observedAt: interval.end + ) + ) + + let updated = cached.updating( + interval: interval, + observation: .gap( + sourceVersion: "0.145.0", + observedAt: interval.end, + reason: "Codex account identity could not be checked" + ) + ) + + XCTAssertEqual(updated.tokens, 100) + XCTAssertEqual(updated.coverage, .low) + XCTAssertEqual( + updated.reason, + "Codex account identity could not be checked" + ) + } + func testNoLocalActivityInAContinuouslyObservedIntervalIsNotApplicable() { let interval = DateInterval( start: Date(timeIntervalSince1970: 1_000), diff --git a/Tests/CodexLimitsTests/RolloutTailSourceTests.swift b/Tests/CodexLimitsTests/RolloutTailSourceTests.swift index 3e5f0f8..c302b65 100644 --- a/Tests/CodexLimitsTests/RolloutTailSourceTests.swift +++ b/Tests/CodexLimitsTests/RolloutTailSourceTests.swift @@ -69,6 +69,26 @@ final class RolloutTailSourceTests: XCTestCase { ) } + func testNegativeTokenCountersAreNotAccepted() throws { + let fixture = try TemporaryRollout() + try fixture.append( + #"{"timestamp":"2026-07-27T10:00:00.000Z","ordinal":1,"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"total_tokens":-1}}}}"# + + "\n" + + #"{"timestamp":"2026-07-27T10:01:00.000Z","ordinal":2,"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":-1,"total_tokens":100}}}}"# + + "\n" + ) + + let batch = try IncrementalRolloutTailSource().read( + fileURL: fixture.url, + cursor: nil, + observedAt: Date(timeIntervalSince1970: 100) + ) + + XCTAssertEqual(batch.records.count, 2) + XCTAssertTrue(batch.records.allSatisfy { $0.tokenUsage == nil }) + XCTAssertEqual(batch.malformedRecordCount, 2) + } + func testLegacyPersistedTokenUsageDoesNotInventComponentPresence() throws { let data = Data( #"{"inputTokens":0,"cachedInputTokens":0,"cacheWriteInputTokens":0,"outputTokens":0,"reasoningOutputTokens":0,"totalTokens":500}"#.utf8 @@ -102,6 +122,7 @@ final class RolloutTailSourceTests: XCTestCase { XCTAssertEqual(batch.records.map(\.ordinal), [0]) XCTAssertEqual(batch.cursor.byteOffset, UInt64(completeLine.utf8.count)) + XCTAssertFalse(batch.hasMoreRecords) XCTAssertEqual( batch.cursor.fileSize, UInt64((completeLine + partialLine).utf8.count) @@ -138,6 +159,101 @@ final class RolloutTailSourceTests: XCTestCase { XCTAssertTrue(replacementBatch.requiresRebuild) } + func testReplacementRebuildsWithoutClaimingContinuityChangedWhenPrefixExceedsBudget() + throws + { + let original = try TemporaryRollout() + let padding = String(repeating: "x", count: 5_000) + let record = + #"{"timestamp":"2026-07-27T10:00:00.000Z","ordinal":0,"type":"session_meta","payload":{"id":"task","padding":"\#(padding)"}}"# + + "\n" + try original.append(record) + let source = IncrementalRolloutTailSource() + let first = try source.read( + fileURL: original.url, + cursor: nil, + observedAt: Date(timeIntervalSince1970: 100) + ) + let replacement = try TemporaryRollout() + try replacement.append(record) + let withoutCheckpoint = RolloutCursor( + fileIdentity: first.cursor.fileIdentity, + sourceGeneration: first.cursor.sourceGeneration, + byteOffset: first.cursor.byteOffset, + fileSize: first.cursor.fileSize, + modificationTime: first.cursor.modificationTime, + lastOrdinal: first.cursor.lastOrdinal, + threadID: first.cursor.threadID, + processedPrefixFingerprint: + first.cursor.processedPrefixFingerprint, + checkpoint: nil, + discardingOversizedRecord: + first.cursor.discardingOversizedRecord + ) + + let rebuilt = try source.read( + fileURL: replacement.url, + cursor: withoutCheckpoint, + observedAt: Date(timeIntervalSince1970: 200), + maximumBytes: 1 + ) + + XCTAssertTrue(rebuilt.requiresRebuild) + XCTAssertFalse(rebuilt.continuityChanged) + XCTAssertTrue(rebuilt.hasMoreRecords) + XCTAssertEqual(rebuilt.cursor.byteOffset, 0) + XCTAssertLessThanOrEqual(rebuilt.bytesRead, 1) + + let completed = try source.read( + fileURL: replacement.url, + cursor: rebuilt.cursor, + observedAt: Date(timeIntervalSince1970: 300) + ) + XCTAssertEqual(completed.records.map(\.ordinal), [0]) + } + + func testGenerationOverflowFailsWithoutTrapping() throws { + let original = try TemporaryRollout() + let record = + #"{"timestamp":"2026-07-27T10:00:00.000Z","ordinal":0,"type":"session_meta","payload":{"id":"task","cli_version":"0.145.0"}}"# + + "\n" + try original.append(record) + let source = IncrementalRolloutTailSource() + let first = try source.read( + fileURL: original.url, + cursor: nil, + observedAt: Date(timeIntervalSince1970: 100) + ) + let exhausted = RolloutCursor( + fileIdentity: first.cursor.fileIdentity, + sourceGeneration: .max, + byteOffset: first.cursor.byteOffset, + fileSize: first.cursor.fileSize, + modificationTime: first.cursor.modificationTime, + lastOrdinal: first.cursor.lastOrdinal, + threadID: first.cursor.threadID, + processedPrefixFingerprint: + first.cursor.processedPrefixFingerprint, + checkpoint: first.cursor.checkpoint, + discardingOversizedRecord: nil + ) + let replacement = try TemporaryRollout() + try replacement.append(record) + + XCTAssertThrowsError( + try source.read( + fileURL: replacement.url, + cursor: exhausted, + observedAt: Date(timeIntervalSince1970: 200) + ) + ) { + XCTAssertEqual( + $0 as? RolloutTailSourceError, + .counterOverflow + ) + } + } + func testTailDoesNotEmitReplayedOrdinalsFromReplacementCopy() throws { let history = #"{"timestamp":"2026-07-27T10:00:00.000Z","ordinal":0,"type":"session_meta","payload":{"id":"task-root","cli_version":"0.145.0"}}"# @@ -259,6 +375,89 @@ final class RolloutTailSourceTests: XCTestCase { XCTAssertEqual(rebuilt.records.last?.totalTokens, 200) } + func testChangedMiddleRecordCannotHideBehindAnUnchangedCheckpoint() throws { + let firstLine = + #"{"timestamp":"2026-07-27T10:00:00.000Z","ordinal":0,"type":"session_meta","payload":{"id":"task-root","cli_version":"0.145.0"}}"# + + "\n" + let middle = + #"{"timestamp":"2026-07-27T10:01:00.000Z","ordinal":1,"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"total_tokens":100}}}}"# + + "\n" + let last = + #"{"timestamp":"2026-07-27T10:02:00.000Z","ordinal":2,"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"total_tokens":300}}}}"# + + "\n" + let original = try TemporaryRollout() + try original.append(firstLine + middle + last) + let source = IncrementalRolloutTailSource() + let first = try source.read( + fileURL: original.url, + cursor: nil, + observedAt: Date(timeIntervalSince1970: 100) + ) + let replacement = try TemporaryRollout() + try replacement.append( + firstLine + + middle.replacingOccurrences( + of: #""total_tokens":100"#, + with: #""total_tokens":200"# + ) + + last + ) + + let rebuilt = try source.read( + fileURL: replacement.url, + cursor: first.cursor, + observedAt: Date(timeIntervalSince1970: 200) + ) + + XCTAssertTrue(rebuilt.requiresRebuild) + XCTAssertEqual( + rebuilt.records.compactMap(\.totalTokens), + [200, 300] + ) + } + + func testSameSizeRewriteChecksMoreThanTheLastRecord() throws { + let firstLine = + #"{"timestamp":"2026-07-27T10:00:00.000Z","ordinal":0,"type":"session_meta","payload":{"id":"task-root","cli_version":"0.145.0"}}"# + + "\n" + let middle = + #"{"timestamp":"2026-07-27T10:01:00.000Z","ordinal":1,"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"total_tokens":100}}}}"# + + "\n" + let last = + #"{"timestamp":"2026-07-27T10:02:00.000Z","ordinal":2,"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"total_tokens":300}}}}"# + + "\n" + let fixture = try TemporaryRollout() + let original = firstLine + middle + last + try fixture.append(original) + let source = IncrementalRolloutTailSource() + let first = try source.read( + fileURL: fixture.url, + cursor: nil, + observedAt: Date(timeIntervalSince1970: 100) + ) + let changed = original.replacingOccurrences( + of: #""total_tokens":100"#, + with: #""total_tokens":200"# + ) + try fixture.replace(with: changed) + try FileManager.default.setAttributes( + [.modificationDate: first.cursor.modificationTime.addingTimeInterval(1)], + ofItemAtPath: fixture.url.path + ) + + let rebuilt = try source.read( + fileURL: fixture.url, + cursor: first.cursor, + observedAt: Date(timeIntervalSince1970: 200) + ) + + XCTAssertTrue(rebuilt.requiresRebuild) + XCTAssertEqual( + rebuilt.records.compactMap(\.totalTokens), + [200, 300] + ) + } + func testChangedReplacementDoesNotReusePreviousTaskIdentity() throws { let original = try TemporaryRollout() try original.append( @@ -381,7 +580,7 @@ final class RolloutTailSourceTests: XCTestCase { XCTAssertEqual(rebuilt.records.last?.totalTokens, 200) } - func testOversizedAcceptedRecordClearsOlderCheckpoint() throws { + func testLargeAcceptedRecordUsesABoundedRawCheckpoint() throws { let fixture = try TemporaryRollout() let padding = String(repeating: "x", count: 5_000) let session = @@ -398,7 +597,10 @@ final class RolloutTailSourceTests: XCTestCase { observedAt: Date(timeIntervalSince1970: 100) ) - XCTAssertNil(first.cursor.checkpoint) + XCTAssertGreaterThan( + try XCTUnwrap(first.cursor.checkpoint).byteLength, + 4_096 + ) let changedContext = originalContext.replacingOccurrences( of: #""model":"gpt-5.6""#, with: #""model":"gpt-5.5""# @@ -414,6 +616,38 @@ final class RolloutTailSourceTests: XCTestCase { XCTAssertEqual(rebuilt.records.last?.model, "gpt-5.5") } + func testAppendAfterLargeAcceptedRecordDoesNotRebuildTheFile() throws { + let fixture = try TemporaryRollout() + let padding = String(repeating: "x", count: 5_000) + try fixture.append( + #"{"timestamp":"2026-07-27T10:01:00.000Z","ordinal":1,"type":"turn_context","payload":{"turn_id":"turn-1","model":"gpt-5.6","effort":"high","padding":"\#(padding)"}}"# + + "\n" + ) + let source = IncrementalRolloutTailSource() + let first = try source.read( + fileURL: fixture.url, + cursor: nil, + observedAt: Date(timeIntervalSince1970: 100) + ) + let appended = + #"{"timestamp":"2026-07-27T10:02:00.000Z","ordinal":2,"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"total_tokens":200}}}}"# + + "\n" + try fixture.append(appended) + + let incremental = try source.read( + fileURL: fixture.url, + cursor: first.cursor, + observedAt: Date(timeIntervalSince1970: 200) + ) + + XCTAssertFalse(incremental.requiresRebuild) + XCTAssertEqual(incremental.records.map(\.ordinal), [2]) + XCTAssertLessThanOrEqual( + incremental.bytesRead, + UInt64(appended.utf8.count + 4_096) + ) + } + func testReplacementContinuesAfterLargePreOrdinalPrefixWithoutReplay() throws { var history = #"{"timestamp":"2026-07-27T10:00:00.000Z","type":"session_meta","payload":{"id":"task-root","cli_version":"0.145.0","history_mode":"inline"}}"# @@ -557,6 +791,27 @@ final class RolloutTailSourceTests: XCTestCase { XCTAssertGreaterThan(batch.cursor.byteOffset, 0) } + func testCompactedRecordSkipsItsUnusedLargePayload() throws { + let fixture = try TemporaryRollout() + let sourceContent = String(repeating: "x", count: 1_000_000) + try fixture.append( + #"{"timestamp":"2026-07-27T10:00:00.000Z","type":"compacted","payload":{"info":"not-a-token-object","replacement_history":"\#(sourceContent)"}}"# + + "\n" + ) + + let batch = try IncrementalRolloutTailSource().read( + fileURL: fixture.url, + cursor: nil, + observedAt: Date(timeIntervalSince1970: 100) + ) + let record = try XCTUnwrap(batch.records.first) + + XCTAssertEqual(record.type, "compacted") + XCTAssertEqual(record.timestamp, "2026-07-27T10:00:00.000Z") + XCTAssertNil(record.tokenUsage) + XCTAssertEqual(batch.malformedRecordCount, 0) + } + func testCompleteMalformedLineIsSkippedWithoutLosingValidRecords() throws { let fixture = try TemporaryRollout() try fixture.append( @@ -567,16 +822,204 @@ final class RolloutTailSourceTests: XCTestCase { + "\n" ) - let batch = try IncrementalRolloutTailSource().read( + let source = IncrementalRolloutTailSource() + let first = try source.read( fileURL: fixture.url, cursor: nil, - observedAt: Date(timeIntervalSince1970: 100) + observedAt: Date(timeIntervalSince1970: 100), + maximumLines: 2 + ) + let second = try source.read( + fileURL: fixture.url, + cursor: first.cursor, + observedAt: Date(timeIntervalSince1970: 200), + maximumLines: 2 ) - XCTAssertEqual(batch.records.map(\.ordinal), [0, 1]) - XCTAssertEqual(batch.unsupportedRecordCount, 1) - XCTAssertEqual(batch.records.last?.totalTokens, 100) - XCTAssertEqual(batch.cursor.lastOrdinal, 1) + XCTAssertEqual(first.records.map(\.ordinal), [0]) + XCTAssertEqual(first.unsupportedRecordCount, 1) + XCTAssertTrue(first.hasMoreRecords) + XCTAssertEqual(second.records.map(\.ordinal), [1]) + XCTAssertEqual(second.records.last?.totalTokens, 100) + XCTAssertEqual(second.cursor.lastOrdinal, 1) + } + + func testOversizedRecordIsSkippedWithoutLosingTheNextRecord() throws { + let fixture = try TemporaryRollout() + let padding = String( + repeating: "x", + count: BoundedJSONLReader.maximumRecordBytes + ) + try fixture.append( + #"{"timestamp":"2026-07-27T10:00:00.000Z","ordinal":1,"type":"event_msg","payload":{"type":"token_count","padding":"\#(padding)","info":{"total_token_usage":{"total_tokens":100}}}}"# + + "\n" + + #"{"timestamp":"2026-07-27T10:01:00.000Z","ordinal":2,"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"total_tokens":200}}}}"# + + "\n" + ) + + let source = IncrementalRolloutTailSource() + let first = try source.read( + fileURL: fixture.url, + cursor: nil, + observedAt: Date(timeIntervalSince1970: 100), + maximumLines: 1 + ) + let second = try source.read( + fileURL: fixture.url, + cursor: first.cursor, + observedAt: Date(timeIntervalSince1970: 200), + maximumLines: 1 + ) + + XCTAssertTrue(first.records.isEmpty) + XCTAssertEqual(first.malformedRecordCount, 1) + XCTAssertTrue(first.hasMoreRecords) + XCTAssertEqual(second.records.map(\.ordinal), [2]) + XCTAssertFalse(second.hasMoreRecords) + XCTAssertEqual(second.cursor.byteOffset, second.cursor.fileSize) + } + + func testOversizedUnterminatedRecordResumesWithoutRescanning() throws { + let fixture = try TemporaryRollout() + try fixture.append(String(repeating: "x", count: 100_000)) + let source = IncrementalRolloutTailSource() + + let first = try source.read( + fileURL: fixture.url, + cursor: nil, + observedAt: Date(timeIntervalSince1970: 100), + maximumBytes: 32_768, + maximumRecordBytes: 16 + ) + var cursor = first.cursor + var readCount = 1 + while cursor.byteOffset < cursor.fileSize { + let next = try source.read( + fileURL: fixture.url, + cursor: cursor, + observedAt: Date( + timeIntervalSince1970: 100 + Double(readCount) + ), + maximumBytes: 32_768, + maximumRecordBytes: 16 + ) + XCTAssertLessThanOrEqual(next.bytesRead, 32_768) + XCTAssertGreaterThan(next.cursor.byteOffset, cursor.byteOffset) + cursor = next.cursor + readCount += 1 + XCTAssertLessThan(readCount, 10) + } + let idle = try source.read( + fileURL: fixture.url, + cursor: cursor, + observedAt: Date(timeIntervalSince1970: 300), + maximumBytes: 32_768, + maximumRecordBytes: 16 + ) + + XCTAssertTrue(first.hasMoreRecords) + XCTAssertTrue(first.cursor.discardingOversizedRecord == true) + XCTAssertGreaterThan(first.cursor.byteOffset, 0) + XCTAssertEqual(cursor.byteOffset, cursor.fileSize) + XCTAssertTrue(cursor.discardingOversizedRecord == true) + XCTAssertEqual(idle.bytesRead, 0) + } + + func testGrowingOversizedUnterminatedRecordContinuesFromItsCursor() + throws + { + let fixture = try TemporaryRollout() + try fixture.append(String(repeating: "x", count: 40_000)) + let source = IncrementalRolloutTailSource() + let first = try source.read( + fileURL: fixture.url, + cursor: nil, + observedAt: Date(timeIntervalSince1970: 100), + maximumBytes: 16_384, + maximumRecordBytes: 16 + ) + try fixture.append(String(repeating: "x", count: 40_000)) + + let second = try source.read( + fileURL: fixture.url, + cursor: first.cursor, + observedAt: Date(timeIntervalSince1970: 200), + maximumBytes: 16_384, + maximumRecordBytes: 16 + ) + + XCTAssertLessThanOrEqual(second.bytesRead, 16_384) + XCTAssertGreaterThan(second.cursor.byteOffset, first.cursor.byteOffset) + XCTAssertFalse(second.requiresRebuild) + } + + func testByteBudgetDefersAValidPartialRecordWithoutReadingPastBudget() + throws + { + let fixture = try TemporaryRollout() + let padding = String(repeating: "x", count: 100_000) + try fixture.append( + #"{"ordinal":1,"type":"session_meta","payload":{"id":"task","padding":"\#(padding)"}}"# + + "\n" + ) + let source = IncrementalRolloutTailSource() + + let batch = try source.read( + fileURL: fixture.url, + cursor: nil, + observedAt: Date(timeIntervalSince1970: 100), + maximumBytes: 32_768, + maximumRecordBytes: 1_048_576 + ) + let completed = try source.read( + fileURL: fixture.url, + cursor: batch.cursor, + observedAt: Date(timeIntervalSince1970: 200), + maximumBytes: 200_000, + maximumRecordBytes: 1_048_576 + ) + + XCTAssertEqual(batch.bytesRead, 32_768) + XCTAssertTrue(batch.hasMoreRecords) + XCTAssertNil(batch.cursor.discardingOversizedRecord) + XCTAssertEqual(batch.cursor.byteOffset, 0) + XCTAssertEqual(batch.malformedRecordCount, 0) + XCTAssertEqual(completed.records.map(\.ordinal), [1]) + } + + func testLineLimitReturnsAResumableBatch() throws { + let fixture = try TemporaryRollout() + let firstRecord = + #"{"ordinal":0,"type":"session_meta","payload":{"id":"task"}}"# + + "\n" + let secondRecord = + #"{"ordinal":1,"type":"event_msg","payload":{"type":"turn_started"}}"# + + "\n" + let thirdRecord = + #"{"ordinal":2,"type":"event_msg","payload":{"type":"turn_started"}}"# + + "\n" + try fixture.append(firstRecord + secondRecord + thirdRecord) + let source = IncrementalRolloutTailSource() + + let first = try source.read( + fileURL: fixture.url, + cursor: nil, + observedAt: Date(timeIntervalSince1970: 100), + maximumLines: 2 + ) + let second = try source.read( + fileURL: fixture.url, + cursor: first.cursor, + observedAt: Date(timeIntervalSince1970: 200), + maximumLines: 2 + ) + + XCTAssertEqual(first.records.map(\.ordinal), [0, 1]) + XCTAssertTrue(first.hasMoreRecords) + XCTAssertLessThan(first.cursor.byteOffset, first.cursor.fileSize) + XCTAssertEqual(second.records.map(\.ordinal), [2]) + XCTAssertFalse(second.hasMoreRecords) + XCTAssertEqual(second.cursor.byteOffset, second.cursor.fileSize) } func testPartialLineIsEmittedOnceAfterItBecomesComplete() throws { diff --git a/Tests/CodexLimitsTests/UsageHistoryTests.swift b/Tests/CodexLimitsTests/UsageHistoryTests.swift index dc7a7a5..43bb09d 100644 --- a/Tests/CodexLimitsTests/UsageHistoryTests.swift +++ b/Tests/CodexLimitsTests/UsageHistoryTests.swift @@ -128,6 +128,47 @@ final class UsageHistoryTests: XCTestCase { XCTAssertTrue(state.samples.isEmpty) } + func testDailyFileRestoreKeepsValidSamplesBesideAnInvalidSample() async throws { + let root = temporaryDirectory() + let valid = UsageSample( + observedAt: Date(timeIntervalSince1970: 1_900_000), + remainingPercent: 80, + resetsAt: Date(timeIntervalSince1970: 2_000_000) + ) + let history = UsageHistory( + localDirectory: root, + installationID: "writer-a" + ) + _ = await history.load() + _ = await history.record(valid) + let file = try XCTUnwrap( + jsonFiles(for: "writer-a", in: root).first + ) + var dailyFile = try XCTUnwrap( + JSONSerialization.jsonObject( + with: Data(contentsOf: file) + ) as? [String: Any] + ) + var samples = try XCTUnwrap( + dailyFile["samples"] as? [[String: Any]] + ) + samples.append([ + "observedAt": 1_900_001, + "remainingPercent": -1, + "resetsAt": 2_000_000 + ]) + dailyFile["samples"] = samples + try JSONSerialization.data(withJSONObject: dailyFile).write(to: file) + + let restored = await UsageHistory( + localDirectory: root, + installationID: "writer-a" + ).load() + + XCTAssertEqual(restored.samples, [valid]) + XCTAssertNil(restored.errorMessage) + } + func testSharedFolderAcceptsTheSameAccountAndRejectsAnotherAccount() async throws { let root = temporaryDirectory() let shared = root.appendingPathComponent("shared", isDirectory: true) @@ -282,6 +323,50 @@ final class UsageHistoryTests: XCTestCase { XCTAssertEqual(sample.resetsAt, Date(timeIntervalSinceReferenceDate: 86_400)) } + func testUsageSampleDecodingLeavesValidationToTheRestoreBoundary() throws { + let invalid = [ + #"{"observedAt":0,"remainingPercent":1e308,"resetsAt":86400}"#, + #"{"observedAt":0,"remainingPercent":-1,"resetsAt":86400}"#, + #"{"observedAt":86401,"remainingPercent":80,"resetsAt":86400}"#, + #"{"observedAt":0,"remainingPercent":80,"resetsAt":86400,"lifetimeTokens":-1}"# + ] + + for json in invalid { + let sample = try JSONDecoder().decode( + UsageSample.self, + from: Data(json.utf8) + ) + XCTAssertFalse(sample.isValid) + } + + let decoder = JSONDecoder() + decoder.nonConformingFloatDecodingStrategy = .convertFromString( + positiveInfinity: "Infinity", + negativeInfinity: "-Infinity", + nan: "NaN" + ) + let sample = try decoder.decode( + UsageSample.self, + from: Data( + #"{"observedAt":0,"remainingPercent":"NaN","resetsAt":86400}"# + .utf8 + ) + ) + XCTAssertFalse(sample.isValid) + } + + func testSpendControlDecodingRejectsInvalidRemainingPercent() throws { + let decoder = JSONDecoder() + let data = Data( + #"{"limit":"50","used":"10","remainingPercent":1e308,"resetsAt":86400}"# + .utf8 + ) + + XCTAssertThrowsError( + try decoder.decode(AccountSpendControlFacts.self, from: data) + ) + } + func testVersionOneHistoryMigratesIntoTheActiveAccountPartition() async throws { let root = temporaryDirectory() let writer = root @@ -1106,6 +1191,49 @@ final class UsageHistoryTests: XCTestCase { XCTAssertNil(state.errorMessage) } + func testAutomaticSyncIsThrottledButExplicitSyncStillRuns() async throws { + let root = temporaryDirectory() + let shared = root.appendingPathComponent("shared", isDirectory: true) + try FileManager.default.createDirectory( + at: shared, + withIntermediateDirectories: true + ) + let first = UsageSample( + observedAt: Date(timeIntervalSince1970: 1_900_000), + remainingPercent: 80, + resetsAt: Date(timeIntervalSince1970: 2_000_000) + ) + let second = UsageSample( + observedAt: Date(timeIntervalSince1970: 1_900_060), + remainingPercent: 79, + resetsAt: Date(timeIntervalSince1970: 2_000_000) + ) + let reader = UsageHistory( + localDirectory: root.appendingPathComponent("reader"), + installationID: "reader" + ) + _ = await reader.load() + _ = await reader.connect(to: shared) + _ = await reader.record(first) + let writer = UsageHistory( + localDirectory: root.appendingPathComponent("writer"), + installationID: "writer" + ) + _ = await writer.load() + _ = await writer.connect(to: shared) + _ = await writer.record(second) + + let throttled = await reader.synchronizeIfDue() + let afterClockRollback = await reader.synchronizeIfDue( + at: Date(timeIntervalSince1970: 0) + ) + let explicit = await reader.synchronize() + + XCTAssertEqual(throttled.samples, [first]) + XCTAssertEqual(afterClockRollback.samples, [first, second]) + XCTAssertEqual(explicit.samples, [first, second]) + } + func testUnsupportedFolderVersionIsNotModified() async throws { let root = temporaryDirectory() let shared = root.appendingPathComponent("shared", isDirectory: true) diff --git a/Tests/CodexLimitsTests/UsageIntelligenceEngineTests.swift b/Tests/CodexLimitsTests/UsageIntelligenceEngineTests.swift index 85be843..01b9c1e 100644 --- a/Tests/CodexLimitsTests/UsageIntelligenceEngineTests.swift +++ b/Tests/CodexLimitsTests/UsageIntelligenceEngineTests.swift @@ -1342,13 +1342,32 @@ final class UsageIntelligenceEngineTests: XCTestCase { localActivityObservation: .continuous( sourceVersion: "test", observedAt: now - ) + ), + localActivityContentRevision: 7 + ) + ) + let reused = UsageIntelligenceEngine.evaluate( + UsageIntelligenceInput( + account: account, + samples: samples, + safetyBuffer: 3, + sourceState: .available, + now: now, + previousStatus: nil, + localActivityFacts: [], + localActivityObservation: .continuous( + sourceVersion: "test", + observedAt: now + ), + localActivityContentRevision: 7, + reusableLocalAggregates: reader.reusableLocalAggregates ) ) XCTAssertEqual(reader.evidence.coverage, .partial) XCTAssertEqual(reader.evidence.confidence, .medium) XCTAssertEqual(reader.evidence.reason, "Workload mix changed") + XCTAssertEqual(reused.evidence, reader.evidence) } func testKnownResetKeepsObservedWindowsSeparate() { @@ -1790,6 +1809,356 @@ final class UsageIntelligenceEngineTests: XCTestCase { ) } + func testStableRevisionReusesLocalAggregatesAcrossLaterAccountRead() throws { + let observedAt = Date(timeIntervalSince1970: 2_000_000) + let account = makeSnapshot(remaining: 80, fetchedAt: observedAt) + let observation = LocalActivityObservation.continuous( + sourceVersion: "0.145.0", + observedAt: observedAt + ) + let first = UsageIntelligenceEngine.evaluate( + UsageIntelligenceInput( + account: account, + samples: [], + safetyBuffer: 3, + sourceState: .available, + now: observedAt, + previousStatus: nil, + localActivityFacts: [ + tokenFact( + tokens: 100, + date: observedAt.addingTimeInterval(-30), + eventID: "token-1" + ) + ], + localActivityObservation: observation, + localActivityContentRevision: 7 + ) + ) + let later = observedAt.addingTimeInterval(60) + let laterObservation = LocalActivityObservation.continuous( + sourceVersion: "0.145.0", + observedAt: later + ) + let laterAccount = UsageSnapshot( + mainLimit: account.mainLimit, + otherLimits: account.otherLimits, + tokenHistory: account.tokenHistory, + emergencyResetCount: account.emergencyResetCount, + bankedResetCountAvailable: account.bankedResetCountAvailable, + bankedResetDetails: account.bankedResetDetails, + fetchedAt: later, + accountFacts: account.accountFacts + ) + + let reused = UsageIntelligenceEngine.evaluate( + UsageIntelligenceInput( + account: laterAccount, + samples: [], + safetyBuffer: 3, + sourceState: .available, + now: later, + previousStatus: nil, + localActivityFacts: [], + localActivityObservation: laterObservation, + localActivityContentRevision: 7, + reusableLocalAggregates: + try XCTUnwrap(first.reusableLocalAggregates) + ) + ) + let rebuilt = UsageIntelligenceEngine.evaluate( + UsageIntelligenceInput( + account: laterAccount, + samples: [], + safetyBuffer: 3, + sourceState: .available, + now: later, + previousStatus: nil, + localActivityFacts: [], + localActivityObservation: laterObservation, + localActivityContentRevision: 8, + reusableLocalAggregates: + try XCTUnwrap(first.reusableLocalAggregates) + ) + ) + + XCTAssertEqual(reused.localTokenActivity.tokens, 100) + XCTAssertEqual(reused.localTokenActivity.interval.end, later) + XCTAssertEqual(rebuilt.localTokenActivity.tokens, 0) + } + + func testStableRevisionReusesHistoryIndexButAppliesNewAccountSample() throws { + let now = Date(timeIntervalSince1970: 2_000_000) + let account = makeSnapshot(remaining: 80, fetchedAt: now) + let window = try XCTUnwrap(account.mainLimit?.window) + let earlier = now.addingTimeInterval(-60) + let earlierAccount = UsageSnapshot( + mainLimit: LimitReading( + limitId: "codex", + name: "Codex", + window: UsageWindow( + remainingPercent: 90, + resetsAt: window.resetsAt, + durationMinutes: window.durationMinutes + ) + ), + otherLimits: [], + tokenHistory: [], + emergencyResetCount: 0, + bankedResetDetails: nil, + fetchedAt: earlier + ) + let startSample = UsageSample( + observedAt: window.startsAt, + remainingPercent: 100, + resetsAt: window.resetsAt, + lifetimeTokens: 1_000 + ) + let earlierSample = UsageSample( + observedAt: earlier, + remainingPercent: 90, + resetsAt: window.resetsAt, + lifetimeTokens: 1_500 + ) + let freshSample = UsageSample( + observedAt: now, + remainingPercent: 80, + resetsAt: window.resetsAt, + lifetimeTokens: 1_600 + ) + let fact = tokenFact( + tokens: 500, + date: earlier.addingTimeInterval(-60), + eventID: "token-1" + ) + let compatibleSources: Set = [ + LocalTokenDefinitionSource( + sourceVersion: "0.145.0", + schemaVersion: "rollout-jsonl-v1" + ) + ] + let first = UsageIntelligenceEngine.evaluate( + UsageIntelligenceInput( + account: earlierAccount, + samples: [startSample, earlierSample], + safetyBuffer: 3, + sourceState: .available, + now: earlier, + previousStatus: nil, + accountPartitionID: "account-a", + localActivityFacts: [fact], + localActivityObservation: .continuous( + sourceVersion: "0.145.0", + observedAt: earlier + ), + localActivityContentRevision: 7, + compatibleTokenSources: compatibleSources + ) + ) + + let refreshed = UsageIntelligenceEngine.evaluate( + UsageIntelligenceInput( + account: account, + samples: [startSample, earlierSample, freshSample], + safetyBuffer: 3, + sourceState: .available, + now: now, + previousStatus: nil, + accountPartitionID: "account-a", + localActivityFacts: [fact], + localActivityHistoryFacts: [fact], + localActivityObservation: .continuous( + sourceVersion: "0.145.0", + observedAt: now + ), + localActivityContentRevision: 7, + reusableLocalAggregates: + try XCTUnwrap(first.reusableLocalAggregates), + compatibleTokenSources: compatibleSources + ) + ) + + XCTAssertEqual( + refreshed.usagePerToken.current?.accountMovementPoints, + 20 + ) + XCTAssertEqual( + refreshed.usagePerToken.current?.accountTokenActivity, + 600 + ) + XCTAssertEqual( + refreshed.usagePerToken.current?.localTokenActivity, + 500 + ) + } + + func testStableRevisionRebuildsAfterTemporaryObservationGap() throws { + let observedAt = Date(timeIntervalSince1970: 2_000_000) + let account = makeSnapshot(remaining: 80, fetchedAt: observedAt) + let fact = tokenFact( + tokens: 100, + date: observedAt.addingTimeInterval(-30), + eventID: "token-1" + ) + let first = UsageIntelligenceEngine.evaluate( + UsageIntelligenceInput( + account: account, + samples: [], + safetyBuffer: 3, + sourceState: .available, + now: observedAt, + previousStatus: nil, + localActivityFacts: [fact], + localActivityObservation: .gap( + sourceVersion: "0.145.0", + observedAt: observedAt, + reason: "Codex account identity could not be checked" + ), + localActivityContentRevision: 7 + ) + ) + let recovered = UsageIntelligenceEngine.evaluate( + UsageIntelligenceInput( + account: account, + samples: [], + safetyBuffer: 3, + sourceState: .available, + now: observedAt, + previousStatus: nil, + localActivityFacts: [fact], + localActivityObservation: .continuous( + sourceVersion: "0.145.0", + observedAt: observedAt + ), + localActivityContentRevision: 7, + reusableLocalAggregates: + try XCTUnwrap(first.reusableLocalAggregates) + ) + ) + + XCTAssertEqual(first.localTokenActivity.coverage, .low) + XCTAssertEqual(recovered.localTokenActivity.coverage, .high) + XCTAssertEqual( + recovered.localTokenActivity.reason, + "Only local activity on this Mac is observed" + ) + } + + func testStableRevisionRebuildsWhenPreviousEndBoundaryBecomesIncluded() throws { + let observedAt = Date(timeIntervalSince1970: 2_000_000) + let account = makeSnapshot(remaining: 80, fetchedAt: observedAt) + let observation = LocalActivityObservation.continuous( + sourceVersion: "0.145.0", + observedAt: observedAt + ) + let boundaryFact = tokenFact( + tokens: 250, + date: observedAt, + eventID: "boundary-token" + ) + let first = UsageIntelligenceEngine.evaluate( + UsageIntelligenceInput( + account: account, + samples: [], + safetyBuffer: 3, + sourceState: .available, + now: observedAt, + previousStatus: nil, + localActivityFacts: [boundaryFact], + localActivityObservation: observation, + localActivityContentRevision: 7 + ) + ) + let later = observedAt.addingTimeInterval(60) + let laterAccount = UsageSnapshot( + mainLimit: account.mainLimit, + otherLimits: account.otherLimits, + tokenHistory: account.tokenHistory, + emergencyResetCount: account.emergencyResetCount, + bankedResetCountAvailable: account.bankedResetCountAvailable, + bankedResetDetails: account.bankedResetDetails, + fetchedAt: later, + accountFacts: account.accountFacts + ) + + let refreshed = UsageIntelligenceEngine.evaluate( + UsageIntelligenceInput( + account: laterAccount, + samples: [], + safetyBuffer: 3, + sourceState: .available, + now: later, + previousStatus: nil, + localActivityFacts: [boundaryFact], + localActivityObservation: .continuous( + sourceVersion: "0.145.0", + observedAt: later + ), + localActivityContentRevision: 7, + reusableLocalAggregates: + try XCTUnwrap(first.reusableLocalAggregates) + ) + ) + + XCTAssertEqual(first.localTokenActivity.tokens, 0) + XCTAssertEqual(refreshed.localTokenActivity.tokens, 250) + } + + func testDisplayDownsamplingKeepsTheFirstAndLastPoint() { + let points = Array(0 ..< 10_000) + + let rendered = downsampledForDisplay(points, limit: 100) + + XCTAssertEqual(rendered.count, 100) + XCTAssertEqual(rendered.first, 0) + XCTAssertEqual(rendered.last, 9_999) + } + + func testNearestDisplayPointUsesChronologicalNeighbors() { + let points = (0 ..< 100_000).map { + LocalTokenActivityPoint( + date: Date(timeIntervalSince1970: TimeInterval($0 * 10)), + tokens: Int64($0) + ) + } + + XCTAssertEqual( + nearestPoint( + in: points, + to: Date(timeIntervalSince1970: 123_456), + date: \.date + )?.tokens, + 12_346 + ) + XCTAssertEqual( + nearestPoint( + in: points, + to: Date(timeIntervalSince1970: -1), + date: \.date + )?.tokens, + 0 + ) + } + + func testNearestDisplayPointIncludesTheRangeEnd() { + let start = Date(timeIntervalSince1970: 100) + let end = Date(timeIntervalSince1970: 200) + let points = [ + LocalTokenActivityPoint(date: start, tokens: 1), + LocalTokenActivityPoint(date: end, tokens: 2) + ] + + XCTAssertEqual( + nearestPoint( + in: points, + to: end, + date: \.date, + within: DateInterval(start: start, end: end) + )?.tokens, + 2 + ) + } + private func makeSnapshot( remaining: Double, fetchedAt: Date, diff --git a/Tests/CodexLimitsTests/UsageMonitorHistoryTests.swift b/Tests/CodexLimitsTests/UsageMonitorHistoryTests.swift index 1a3b4a2..c11db76 100644 --- a/Tests/CodexLimitsTests/UsageMonitorHistoryTests.swift +++ b/Tests/CodexLimitsTests/UsageMonitorHistoryTests.swift @@ -4,6 +4,166 @@ import XCTest @MainActor final class UsageMonitorHistoryTests: XCTestCase { + func testSafetyBufferPolicyNormalizesInvalidValues() { + XCTAssertEqual(SafetyBufferPolicy.normalized(nil), 3) + XCTAssertEqual(SafetyBufferPolicy.normalized(.nan), 3) + XCTAssertEqual(SafetyBufferPolicy.normalized(-.infinity), 3) + XCTAssertEqual(SafetyBufferPolicy.normalized(0), 1) + XCTAssertEqual(SafetyBufferPolicy.normalized(1), 1) + XCTAssertEqual(SafetyBufferPolicy.normalized(10), 10) + XCTAssertEqual(SafetyBufferPolicy.normalized(11), 10) + XCTAssertEqual(SafetyBufferPolicy.normalized(.infinity), 3) + } + + func testMonitorRepairsInvalidStoredSafetyBuffer() throws { + for (raw, expected) in [(Double.nan, 3.0), (1e308, 10.0)] { + let suiteName = "UsageMonitorHistoryTests-\(UUID().uuidString)" + let defaults = try XCTUnwrap(UserDefaults(suiteName: suiteName)) + defer { defaults.removePersistentDomain(forName: suiteName) } + defaults.set(raw, forKey: UsageMonitor.safetyBufferKey) + let monitor = UsageMonitor( + defaults: defaults, + historyDirectory: temporaryDirectory(), + startsAutomatically: false + ) + + XCTAssertEqual( + defaults.double(forKey: UsageMonitor.safetyBufferKey), + expected + ) + monitor.updateSafetyBuffer(.infinity) + XCTAssertEqual( + defaults.double(forKey: UsageMonitor.safetyBufferKey), + 3 + ) + } + } + + func testStoredStateRestoreDropsOnlyInvalidSamples() throws { + let suiteName = "UsageMonitorHistoryTests-\(UUID().uuidString)" + let defaults = try XCTUnwrap(UserDefaults(suiteName: suiteName)) + defer { defaults.removePersistentDomain(forName: suiteName) } + let valid = UsageSample( + observedAt: Date(timeIntervalSince1970: 1_900_000), + remainingPercent: 80, + resetsAt: Date(timeIntervalSince1970: 2_000_000) + ) + let invalid = UsageSample( + observedAt: Date(timeIntervalSince1970: 1_900_001), + remainingPercent: -1, + resetsAt: Date(timeIntervalSince1970: 2_000_000) + ) + defaults.set( + try JSONEncoder().encode( + StoredStateFixture( + snapshot: nil, + samples: [valid, invalid], + previousStatus: nil + ) + ), + forKey: "usageState" + ) + + let monitor = UsageMonitor( + defaults: defaults, + historyDirectory: temporaryDirectory(), + startsAutomatically: false + ) + + XCTAssertEqual(monitor.samples, [valid]) + } + + func testStoredAccountEvaluationDoesNotBlockMainActor() async throws { + let suiteName = "UsageMonitorHistoryTests-\(UUID().uuidString)" + let defaults = try XCTUnwrap(UserDefaults(suiteName: suiteName)) + defer { defaults.removePersistentDomain(forName: suiteName) } + let fetchedAt = Date(timeIntervalSince1970: 1_900_000) + defaults.set( + try JSONEncoder().encode( + StoredStateFixture( + snapshot: UsageSnapshot( + mainLimit: LimitReading( + limitId: "codex", + name: "Codex", + window: UsageWindow( + remainingPercent: 64, + resetsAt: fetchedAt.addingTimeInterval(86_400), + durationMinutes: 10_080 + ) + ), + otherLimits: [], + tokenHistory: [], + emergencyResetCount: 0, + fetchedAt: fetchedAt + ), + samples: [], + previousStatus: nil + ) + ), + forKey: "usageState" + ) + let evaluationStarted = expectation( + description: "stored account evaluation started" + ) + let evaluator = BlockingUsageEvaluator(started: evaluationStarted) + + let monitor = UsageMonitor( + defaults: defaults, + historyDirectory: temporaryDirectory(), + startsAutomatically: false, + evaluateUsage: { evaluator.evaluate($0) } + ) + + await fulfillment(of: [evaluationStarted], timeout: 2) + let mainActorResponded = expectation( + description: "main actor stayed responsive" + ) + Task { @MainActor in + mainActorResponded.fulfill() + } + await fulfillment(of: [mainActorResponded], timeout: 2) + evaluator.release() + let deadline = ContinuousClock.now + .seconds(2) + while monitor.readerSnapshot.menuBarText != "64%", + ContinuousClock.now < deadline { + await Task.yield() + } + XCTAssertEqual(monitor.readerSnapshot.menuBarText, "64%") + } + + func testStoredStateRestoreDropsSnapshotWithAnUnboundedDate() throws { + let suiteName = "UsageMonitorHistoryTests-\(UUID().uuidString)" + let defaults = try XCTUnwrap(UserDefaults(suiteName: suiteName)) + defer { defaults.removePersistentDomain(forName: suiteName) } + let fetchedAt = Date( + timeIntervalSinceReferenceDate: -Double.greatestFiniteMagnitude + ) + defaults.set( + try JSONEncoder().encode( + StoredStateFixture( + snapshot: UsageSnapshot( + mainLimit: nil, + otherLimits: [], + tokenHistory: [], + emergencyResetCount: 0, + fetchedAt: fetchedAt + ), + samples: [], + previousStatus: nil + ) + ), + forKey: "usageState" + ) + + let monitor = UsageMonitor( + defaults: defaults, + historyDirectory: temporaryDirectory(), + startsAutomatically: false + ) + + XCTAssertNil(monitor.readerSnapshot.account) + } + func testDeleteAnalyticsHistoryPreservesPreferencesAndDoesNotRestoreLegacySamples() async throws { let suiteName = "UsageMonitorHistoryTests-\(UUID().uuidString)" let defaults = try XCTUnwrap(UserDefaults(suiteName: suiteName)) @@ -427,6 +587,101 @@ final class UsageMonitorHistoryTests: XCTestCase { ) } + func testBlockedEvaluationDoesNotFreezeMainActorOrPublishAStaleSnapshot() async throws { + let suiteName = "UsageMonitorHistoryTests-\(UUID().uuidString)" + let defaults = try XCTUnwrap(UserDefaults(suiteName: suiteName)) + defer { defaults.removePersistentDomain(forName: suiteName) } + let root = temporaryDirectory() + let evaluationStarted = expectation( + description: "first account evaluation started" + ) + let evaluator = BlockingUsageEvaluator(started: evaluationStarted) + let source = FetchSequence([ + makeFetchResult( + identity: "user@example.com", + fetchedAt: Date(timeIntervalSince1970: 1_900_000), + remaining: 80 + ) + ]) + let monitor = UsageMonitor( + defaults: defaults, + historyDirectory: root, + startsAutomatically: false, + fetchUsage: { try await source.next() }, + evaluateUsage: { evaluator.evaluate($0) } + ) + + let refresh = Task { await monitor.refresh() } + await fulfillment(of: [evaluationStarted], timeout: 2) + + let mainActorResponded = expectation( + description: "main actor stayed responsive" + ) + Task { @MainActor in + mainActorResponded.fulfill() + } + await fulfillment(of: [mainActorResponded], timeout: 2) + monitor.updateSafetyBuffer(7) + + evaluator.release() + await refresh.value + let deadline = ContinuousClock.now + .seconds(2) + while monitor.readerSnapshot.chart.target.last?.remaining != 7, + ContinuousClock.now < deadline { + await Task.yield() + } + XCTAssertEqual( + monitor.readerSnapshot.chart.target.last?.remaining, + 7 + ) + } + + func testPreferenceChangeSupersedesAPendingEvaluation() async throws { + let suiteName = "UsageMonitorHistoryTests-\(UUID().uuidString)" + let defaults = try XCTUnwrap(UserDefaults(suiteName: suiteName)) + defer { defaults.removePersistentDomain(forName: suiteName) } + let evaluationStarted = expectation( + description: "account evaluation started" + ) + let evaluator = BlockingUsageEvaluator(started: evaluationStarted) + let source = FetchSequence([ + makeFetchResult( + identity: "user@example.com", + fetchedAt: Date(timeIntervalSince1970: 1_900_000), + remaining: 80 + ) + ]) + let monitor = UsageMonitor( + defaults: defaults, + historyDirectory: temporaryDirectory(), + startsAutomatically: false, + fetchUsage: { try await source.next() }, + evaluateUsage: { evaluator.evaluate($0) } + ) + let refresh = Task { await monitor.refresh() } + await fulfillment(of: [evaluationStarted], timeout: 2) + var exploration = AnalyticsExplorationState.initial + exploration.section = .insights + exploration.timeRange = .threeDays + + monitor.analyticsPreferencesDidChange( + exploration: exploration, + dispositions: [:] + ) + evaluator.release() + await refresh.value + let deadline = ContinuousClock.now + .seconds(2) + while ( + evaluator.callCount < 2 + || monitor.readerSnapshot.menuBarText != "80%" + ), ContinuousClock.now < deadline { + await Task.yield() + } + + XCTAssertEqual(evaluator.lastExploration, exploration) + XCTAssertEqual(monitor.readerSnapshot.menuBarText, "80%") + } + func testFailedRefreshRetainsTheLastReaderSnapshot() async throws { let suiteName = "UsageMonitorHistoryTests-\(UUID().uuidString)" let defaults = try XCTUnwrap(UserDefaults(suiteName: suiteName)) @@ -572,6 +827,86 @@ final class UsageMonitorHistoryTests: XCTestCase { ) } + func testMonitorFinishesABoundedLocalImportWithoutAnotherAccountRead() + async throws + { + let suiteName = "UsageMonitorHistoryTests-\(UUID().uuidString)" + let defaults = try XCTUnwrap(UserDefaults(suiteName: suiteName)) + defer { defaults.removePersistentDomain(forName: suiteName) } + let root = temporaryDirectory() + let rolloutDirectory = root.appendingPathComponent( + "rollouts/1970/01/22", + isDirectory: true + ) + try FileManager.default.createDirectory( + at: rolloutDirectory, + withIntermediateDirectories: true + ) + let fetchedAt = Date(timeIntervalSince1970: 1_850_000) + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [ + .withInternetDateTime, + .withFractionalSeconds + ] + let eventTimestamp = formatter.string( + from: fetchedAt.addingTimeInterval(-1) + ) + var rollout = + #"{"timestamp":"\#(eventTimestamp)","ordinal":0,"type":"session_meta","payload":{"id":"task","cli_version":"0.145.0"}}"# + + "\n" + for ordinal in 1 ... 10_100 { + rollout += + #"{"timestamp":"\#(eventTimestamp)","ordinal":\#(ordinal),"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"total_tokens":\#(ordinal * 100)}}}}"# + + "\n" + } + let rolloutURL = rolloutDirectory.appendingPathComponent( + "rollout-test.jsonl" + ) + let collector = LocalActivityCollector( + rootDirectory: root.appendingPathComponent("rollouts"), + stateDirectory: root.appendingPathComponent("collector-state") + ) + let fetches = FetchSequence([ + makeFetchResult( + identity: "user@example.com", + fetchedAt: Date(timeIntervalSince1970: 1_700_000), + remaining: 90 + ), + makeFetchResult( + identity: "user@example.com", + fetchedAt: fetchedAt, + remaining: 80 + ) + ]) + let monitor = UsageMonitor( + defaults: defaults, + historyDirectory: root.appendingPathComponent("history"), + startsAutomatically: false, + localActivityCollector: collector, + fetchUsage: { try await fetches.next() } + ) + + await monitor.refresh() + try Data(rollout.utf8).write(to: rolloutURL) + await monitor.refresh() + for _ in 0 ..< 100 { + if await collector.hasPendingImport() == false, + monitor.readerSnapshot.localTokenActivity.tokens == 1_009_900 { + break + } + try await Task.sleep(for: .milliseconds(50)) + } + + let pendingAfterContinuation = await collector.hasPendingImport() + XCTAssertFalse(pendingAfterContinuation) + XCTAssertEqual( + monitor.readerSnapshot.localTokenActivity.tokens, + 1_009_900 + ) + let accountReadCount = await fetches.callCount + XCTAssertEqual(accountReadCount, 2) + } + func testMissingWeeklyRefreshClearsWeeklyOutputsAndKeepsOtherLimits() async throws { let suiteName = "UsageMonitorHistoryTests-\(UUID().uuidString)" let defaults = try XCTUnwrap(UserDefaults(suiteName: suiteName)) @@ -975,6 +1310,71 @@ final class UsageMonitorHistoryTests: XCTestCase { XCTAssertEqual(restarted.samples.map(\.remainingPercent), [79, 78]) } + func testManualRefreshRunsSyncInsideTheAutomaticThrottleWindow() + async throws + { + let suiteName = "UsageMonitorHistoryTests-\(UUID().uuidString)" + let defaults = try XCTUnwrap(UserDefaults(suiteName: suiteName)) + defer { defaults.removePersistentDomain(forName: suiteName) } + let root = temporaryDirectory() + let local = root.appendingPathComponent("local", isDirectory: true) + let shared = root.appendingPathComponent("shared", isDirectory: true) + try FileManager.default.createDirectory( + at: shared, + withIntermediateDirectories: true + ) + let source = FetchSequence([ + makeFetchResult( + identity: "user@example.com", + fetchedAt: Date(timeIntervalSince1970: 1_900_000), + remaining: 80 + ), + makeFetchResult( + identity: "user@example.com", + fetchedAt: Date(timeIntervalSince1970: 1_900_060), + remaining: 79 + ) + ]) + let monitor = UsageMonitor( + defaults: defaults, + historyDirectory: local, + startsAutomatically: false, + fetchUsage: { try await source.next() } + ) + await monitor.connectHistoryFolder(shared) + XCTAssertNil(monitor.syncErrorMessage, monitor.syncErrorMessage ?? "") + XCTAssertEqual(monitor.syncFolderName, "shared") + let partition = try JSONDecoder().decode( + AccountHistoryPartition.self, + from: XCTUnwrap(defaults.data(forKey: "historyAccountPartition")) + ) + let remoteWriter = UsageHistory( + localDirectory: root.appendingPathComponent( + "remote", + isDirectory: true + ), + installationID: "remote", + partition: partition + ) + _ = await remoteWriter.load() + let connected = await remoteWriter.connect( + to: shared, + accountIdentity: "user@example.com" + ) + XCTAssertNil(connected.errorMessage, connected.errorMessage ?? "") + XCTAssertEqual(connected.folderName, "shared") + let recorded = await remoteWriter.record(UsageSample( + observedAt: Date(timeIntervalSince1970: 1_900_030), + remainingPercent: 78, + resetsAt: Date(timeIntervalSince1970: 2_000_000) + )) + XCTAssertEqual(recorded.samples.map(\.remainingPercent), [78]) + + await monitor.refresh() + + XCTAssertEqual(monitor.samples.map(\.remainingPercent), [78, 79]) + } + func testLegacySamplesMigrateAfterTheAccountIsObserved() async throws { let suiteName = "UsageMonitorHistoryTests-\(UUID().uuidString)" let defaults = try XCTUnwrap(UserDefaults(suiteName: suiteName)) @@ -1444,17 +1844,21 @@ final class UsageMonitorHistoryTests: XCTestCase { private actor FetchSequence { private var results: [CodexFetchResult] + private var calls = 0 init(_ results: [CodexFetchResult]) { self.results = results } func next() throws -> CodexFetchResult { + calls += 1 guard !results.isEmpty else { throw CodexClientError.invalidResponse } return results.removeFirst() } + + var callCount: Int { calls } } private actor DelayedFetchSource { @@ -1474,6 +1878,53 @@ private actor DelayedFetchSource { } } +private final class BlockingUsageEvaluator: @unchecked Sendable { + private let lock = NSLock() + private let gate = DispatchSemaphore(value: 0) + private let started: XCTestExpectation + private var blockedFirstAccountEvaluation = false + private var inputs: [UsageIntelligenceInput] = [] + + init(started: XCTestExpectation) { + self.started = started + } + + func evaluate( + _ input: UsageIntelligenceInput + ) -> UsageReaderSnapshot { + lock.lock() + inputs.append(input) + let shouldBlock = input.account != nil + && !blockedFirstAccountEvaluation + if shouldBlock { + blockedFirstAccountEvaluation = true + } + lock.unlock() + + if shouldBlock { + started.fulfill() + gate.wait() + } + return UsageIntelligenceEngine.evaluate(input) + } + + func release() { + gate.signal() + } + + var callCount: Int { + lock.lock() + defer { lock.unlock() } + return inputs.count + } + + var lastExploration: AnalyticsExplorationState? { + lock.lock() + defer { lock.unlock() } + return inputs.last?.analyticsExploration + } +} + private struct StoredStateFixture: Codable { let snapshot: UsageSnapshot? let samples: [UsageSample] diff --git a/Tests/CodexLimitsTests/UsageReceiptTests.swift b/Tests/CodexLimitsTests/UsageReceiptTests.swift index f8eea11..3d28846 100644 --- a/Tests/CodexLimitsTests/UsageReceiptTests.swift +++ b/Tests/CodexLimitsTests/UsageReceiptTests.swift @@ -3,6 +3,57 @@ import XCTest @testable import CodexLimits final class UsageReceiptTests: XCTestCase { + func testDecodedTokenTotalOverflowIsUnavailable() throws { + let interval = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 2_000) + ) + let facts = [ + tokenFact( + eventID: "max", + date: Date(timeIntervalSince1970: 1_100), + tokens: .max, + taskID: "root" + ), + tokenFact( + eventID: "one-more", + date: Date(timeIntervalSince1970: 1_200), + tokens: 1, + taskID: "root" + ) + ] + let decodedFacts = try JSONDecoder().decode( + [LocalActivityFact].self, + from: JSONEncoder().encode(facts) + ) + + let snapshot = UsageReceiptAggregator.evaluate( + facts: decodedFacts, + projections: [ + projection(taskID: "root", project: "atlas") + ], + interval: interval, + observation: continuousObservation(for: interval) + ) + let slice = snapshot.slice(in: interval, filters: .all) + let localTokens = snapshot.localTokenSlice( + in: interval, + filters: .all + ) + + XCTAssertEqual(slice.coverage, .unavailable) + XCTAssertEqual(slice.receiptCoverage, .unavailable) + XCTAssertEqual(slice.reason, "Local token total is invalid") + XCTAssertEqual(slice.receiptReason, "Local token total is invalid") + XCTAssertEqual(slice.totalTokens, 0) + XCTAssertTrue(slice.receipts.isEmpty) + XCTAssertTrue(slice.points.isEmpty) + XCTAssertEqual(localTokens.coverage, .unavailable) + XCTAssertEqual(localTokens.reason, "Local token total is invalid") + XCTAssertEqual(localTokens.tokens, 0) + XCTAssertTrue(localTokens.points.isEmpty) + } + func testEnginePublishesReceiptsInTheReaderSnapshot() { let start = Date(timeIntervalSince1970: 1_000) let fetchedAt = Date(timeIntervalSince1970: 2_000) @@ -114,6 +165,122 @@ final class UsageReceiptTests: XCTestCase { XCTAssertEqual(slice.receipts.map(\.tokens), [140, 60]) } + func testOverviewMatchesReceiptTotalsWithoutBuildingTaskTrees() { + let interval = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 2_000) + ) + let snapshot = UsageReceiptAggregator.evaluate( + facts: [ + tokenFact( + eventID: "root-1-a", + date: Date(timeIntervalSince1970: 1_100), + tokens: 100, + taskID: "root-1" + ), + tokenFact( + eventID: "child-1-a", + date: Date(timeIntervalSince1970: 1_200), + tokens: 40, + taskID: "child-1" + ), + tokenFact( + eventID: "root-2-a", + date: Date(timeIntervalSince1970: 1_300), + tokens: 60, + taskID: "root-2" + ), + tokenFact( + eventID: "unattributed", + date: Date(timeIntervalSince1970: 1_400), + tokens: 25, + taskID: nil + ) + ], + projections: [ + projection(taskID: "root-1", project: "atlas"), + projection( + taskID: "child-1", + parentTaskID: "root-1", + project: "atlas" + ), + projection(taskID: "root-2", project: "atlas") + ], + interval: interval, + observation: continuousObservation(for: interval) + ) + + let overview = snapshot.overview(in: interval, filters: .all) + let detailed = snapshot.slice(in: interval, filters: .all) + + XCTAssertEqual(overview.totalTokens, detailed.totalTokens) + XCTAssertEqual(overview.unattributedTokens, detailed.unattributedTokens) + XCTAssertEqual(overview.coverage, detailed.coverage) + XCTAssertEqual(overview.receiptCoverage, detailed.receiptCoverage) + XCTAssertEqual( + overview.receipts.map(\.rootTaskID), + detailed.receipts.map(\.rootTaskID) + ) + XCTAssertEqual( + overview.receipts.map(\.projectLabel), + detailed.receipts.map(\.projectLabel) + ) + XCTAssertEqual( + overview.receipts.map(\.tokens), + detailed.receipts.map(\.tokens) + ) + XCTAssertEqual(overview.receipts.map(\.taskCount), [2, 1]) + } + + func testReceiptLoadsOnlyTheRequestedRootTask() { + let interval = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 2_000) + ) + let snapshot = UsageReceiptAggregator.evaluate( + facts: [ + tokenFact( + eventID: "root-1-a", + date: Date(timeIntervalSince1970: 1_100), + tokens: 100, + taskID: "root-1" + ), + tokenFact( + eventID: "root-2-a", + date: Date(timeIntervalSince1970: 1_300), + tokens: 60, + taskID: "root-2" + ), + tokenFact( + eventID: "unattributed", + date: Date(timeIntervalSince1970: 1_400), + tokens: 25, + taskID: nil + ) + ], + projections: [ + projection(taskID: "root-1", project: "atlas"), + projection(taskID: "root-2", project: "atlas") + ], + interval: interval, + observation: continuousObservation(for: interval) + ) + + let receipt = snapshot.receipt( + rootTaskID: "root-2", + in: interval, + filters: .all + ) + + XCTAssertEqual(receipt?.rootTaskID, "root-2") + XCTAssertEqual(receipt?.tokens, 60) + XCTAssertEqual( + receipt?.reason, + snapshot.overview(in: interval, filters: .all) + .receipts.first { $0.rootTaskID == "root-2" }?.reason + ) + } + func testSelectedRangeKeepsOneTaskAcrossDaysAndDeduplicatesReplay() { let interval = DateInterval( start: Date(timeIntervalSince1970: 1_000), @@ -257,10 +424,23 @@ final class UsageReceiptTests: XCTestCase { reasoning: "high" ) ) + let localTokens = snapshot.localTokenSlice( + in: interval, + filters: WorkspaceFilters( + projectID: "atlas", + taskTreeID: "atlas-task", + model: "gpt-5.6-sol", + reasoning: "high" + ) + ) XCTAssertEqual(slice.receipts.map(\.rootTaskID), ["atlas-task"]) XCTAssertEqual(slice.totalTokens, 90) XCTAssertEqual(slice.points.last?.tokens, 90) + XCTAssertEqual(localTokens.tokens, slice.totalTokens) + XCTAssertEqual(localTokens.points, slice.points) + XCTAssertEqual(localTokens.coverage, slice.coverage) + XCTAssertEqual(localTokens.reason, slice.reason) XCTAssertEqual( snapshot.filterOptions(in: interval), UsageReceiptFilterOptions( @@ -1070,6 +1250,7 @@ final class UsageReceiptTests: XCTestCase { model: "gpt-5.6-sol", reasoning: "high", turnID: "turn-1", + modelContextWindow: 272_000, tokenDelta: LocalTokenUsage( inputTokens: 300, cachedInputTokens: 100, @@ -1077,23 +1258,15 @@ final class UsageReceiptTests: XCTestCase { outputTokens: 60, reasoningOutputTokens: 20, totalTokens: 360 - ) - ), - diagnosticFact( - key: .context, - value: .tokens( - LocalTokenUsage( - inputTokens: 800, - cachedInputTokens: 300, - cacheWriteInputTokens: 15, - outputTokens: 140, - reasoningOutputTokens: 60, - totalTokens: 940 - ) ), - eventID: "context", - date: interval.start.addingTimeInterval(21), - context: turnContext + contextUsage: LocalTokenUsage( + inputTokens: 800, + cachedInputTokens: 300, + cacheWriteInputTokens: 15, + outputTokens: 140, + reasoningOutputTokens: 60, + totalTokens: 940 + ) ), diagnosticFact( key: .compaction, @@ -1532,6 +1705,40 @@ final class UsageReceiptTests: XCTestCase { ) } + func testAggregatorDoesNotRetainTokenFactsBeforeItsInterval() { + let interval = testInterval() + let oldDate = interval.start.addingTimeInterval(-60) + let currentDate = interval.start.addingTimeInterval(60) + let snapshot = UsageReceiptAggregator.evaluate( + facts: [ + tokenFact( + eventID: "old", + date: oldDate, + tokens: 40, + taskID: "root" + ), + tokenFact( + eventID: "current", + date: currentDate, + tokens: 60, + taskID: "root" + ) + ], + projections: [projection(taskID: "root", project: "atlas")], + interval: interval, + observation: continuousObservation(for: interval) + ) + let wide = DateInterval( + start: oldDate.addingTimeInterval(-1), + end: interval.end + ) + + XCTAssertEqual( + snapshot.slice(in: wide, filters: .all).totalTokens, + 60 + ) + } + private func tokenFact( eventID: String, date: Date, @@ -1540,9 +1747,11 @@ final class UsageReceiptTests: XCTestCase { model: String? = nil, reasoning: String? = nil, turnID: String? = nil, + modelContextWindow: Int64? = nil, agent: LocalAgentIdentity? = nil, reason: String? = nil, - tokenDelta: LocalTokenUsage? = nil + tokenDelta: LocalTokenUsage? = nil, + contextUsage: LocalTokenUsage? = nil ) -> LocalActivityFact { LocalActivityFact( key: .token, @@ -1559,9 +1768,11 @@ final class UsageReceiptTests: XCTestCase { turnID: turnID, agent: agent, effectiveModel: model, - reasoning: reasoning + reasoning: reasoning, + modelContextWindow: modelContextWindow ), - tokenDelta: tokenDelta + tokenDelta: tokenDelta, + contextUsage: contextUsage ) } From 49b8f5a3f369d72a9748ef91515476c4d9512b59 Mon Sep 17 00:00:00 2001 From: thrr87 <193831865+thrr87@users.noreply.github.com> Date: Thu, 30 Jul 2026 09:12:13 +0200 Subject: [PATCH 4/5] perf: load local analytics on demand --- Sources/CodexLimits/AnalyticsWorkspace.swift | 4 + .../CodexLimits/LocalActivityCollector.swift | 19 +- Sources/CodexLimits/MenuContentView.swift | 13 +- Sources/CodexLimits/UsageMonitor.swift | 123 +++++++++---- .../AnalyticsWorkspaceTests.swift | 11 ++ .../LocalActivityCollectorTests.swift | 107 +++++++++++ .../UsageMonitorHistoryTests.swift | 173 +++++++++++++++++- 7 files changed, 413 insertions(+), 37 deletions(-) diff --git a/Sources/CodexLimits/AnalyticsWorkspace.swift b/Sources/CodexLimits/AnalyticsWorkspace.swift index 5afc9f5..629b580 100644 --- a/Sources/CodexLimits/AnalyticsWorkspace.swift +++ b/Sources/CodexLimits/AnalyticsWorkspace.swift @@ -101,6 +101,10 @@ struct AnalyticsExplorationState: Codable, Equatable, Sendable { pinnedUsageBaselineID: nil, pinnedUsageBaselineAccountPartitionID: nil ) + + var usesLocalActivity: Bool { + section != .graphs || graph != .usageRemaining + } } @MainActor diff --git a/Sources/CodexLimits/LocalActivityCollector.swift b/Sources/CodexLimits/LocalActivityCollector.swift index 4b81c80..09aabd3 100644 --- a/Sources/CodexLimits/LocalActivityCollector.swift +++ b/Sources/CodexLimits/LocalActivityCollector.swift @@ -816,6 +816,21 @@ actor LocalActivityCollector { importContinuationPending } + func releaseCachedFacts() { + refreshGeneration = nextRevision(after: refreshGeneration) + guard stateURL != nil else { return } + if !changedPaths.isEmpty, !persist() { + return + } + guard changedPaths.isEmpty, + pendingFactWrites.isEmpty else { + return + } + restartPartialFactRestores() + clearPublishedContent() + unloadPersistedFacts(activePaths: []) + } + func deleteHistory(at deletedAt: Date = Date()) throws { historyCutoff = deletedAt historyDeletionPending = stateDirectory != nil @@ -1256,9 +1271,9 @@ actor LocalActivityCollector { guard let directory = stateURL else { return } for path in Array(files.keys) { guard var state = files[path], - state.factsLoaded, !changedPaths.contains(path), - pendingFactWrites[path] == nil else { + pendingFactWrites[path] == nil, + state.factsLoaded || !activePaths.contains(path) else { continue } let file = factsURL( diff --git a/Sources/CodexLimits/MenuContentView.swift b/Sources/CodexLimits/MenuContentView.swift index 203a8c9..b9a1267 100644 --- a/Sources/CodexLimits/MenuContentView.swift +++ b/Sources/CodexLimits/MenuContentView.swift @@ -75,7 +75,16 @@ struct MenuContentView: View { .padding(.vertical, 12) } .frame(width: layout.width, height: layout.height) - .task { await monitor.refresh(forceHistorySync: false) } + .task(id: workspace.state.usesLocalActivity) { + await monitor.setLocalAnalyticsVisible( + workspace.state.usesLocalActivity + ) + } + .onDisappear { + Task { + await monitor.setLocalAnalyticsVisible(false) + } + } .environment(\.locale, Locale(identifier: "en_US")) } @@ -2221,7 +2230,7 @@ private struct UsageRemainingChart: View { .chartOverlay { proxy in chartOverlay(proxy: proxy) } - .frame(height: 300) + .frame(height: 240) .accessibilityElement(children: .ignore) .accessibilityLabel("Usage remaining") .accessibilityValue( diff --git a/Sources/CodexLimits/UsageMonitor.swift b/Sources/CodexLimits/UsageMonitor.swift index ec6458a..466a32f 100644 --- a/Sources/CodexLimits/UsageMonitor.swift +++ b/Sources/CodexLimits/UsageMonitor.swift @@ -80,6 +80,8 @@ final class UsageMonitor: ObservableObject { private var evaluationTask: Task? private var localImportGeneration: UInt64 = 0 private var localImportTask: Task? + private var localAnalyticsVisible = false + private var localAnalyticsNeedsLoad = false convenience init() { self.init( @@ -200,9 +202,7 @@ final class UsageMonitor: ObservableObject { .autoconnect() .sink { [weak self] _ in Task { - @MainActor in await self?.refresh( - forceHistorySync: false - ) + @MainActor in await self?.automaticRefresh() } } .store(in: &cancellables) @@ -211,21 +211,31 @@ final class UsageMonitor: ObservableObject { .publisher(for: NSWorkspace.didWakeNotification) .sink { [weak self] _ in Task { - @MainActor in await self?.refresh( - forceHistorySync: false - ) + @MainActor in await self?.automaticRefresh() } } .store(in: &cancellables) - await refresh(forceHistorySync: false) + await automaticRefresh() + } + + func automaticRefresh() async { + await refresh( + forceHistorySync: false, + includeLocalActivity: false + ) } - func refresh(forceHistorySync: Bool = true) async { + func refresh( + forceHistorySync: Bool = true, + includeLocalActivity: Bool = true + ) async { guard !isRefreshing else { return } isRefreshing = true defer { isRefreshing = false } - cancelLocalImport() + if includeLocalActivity { + cancelLocalImport() + } await restoreHistoryIfAvailable() let fetchTask = Task { try await fetchUsage() } @@ -239,14 +249,14 @@ final class UsageMonitor: ObservableObject { ) accountSnapshot = result.snapshot sourceState = .available - await localActivityCollector?.selectPartition( - historyPartition.id - ) - await refreshLocalActivity( - for: result.snapshot, - observedAt: result.snapshot.fetchedAt, - identityVerified: false - ) + if localAnalyticsVisible, + includeLocalActivity || localAnalyticsNeedsLoad { + await refreshLocalActivity( + for: result.snapshot, + observedAt: result.snapshot.fetchedAt, + identityVerified: false + ) + } let published = await recalculate( now: result.snapshot.fetchedAt ) @@ -297,10 +307,13 @@ final class UsageMonitor: ObservableObject { } accountSnapshot = newSnapshot sourceState = .available - await refreshLocalActivity( - for: newSnapshot, - observedAt: newSnapshot.fetchedAt - ) + if localAnalyticsVisible, + includeLocalActivity || localAnalyticsNeedsLoad { + await refreshLocalActivity( + for: newSnapshot, + observedAt: newSnapshot.fetchedAt + ) + } let published = await recalculate(now: newSnapshot.fetchedAt) persist() if published { @@ -314,10 +327,9 @@ final class UsageMonitor: ObservableObject { (error as? CodexClientError)?.localizedDescription ?? "Couldn’t read Codex usage. Try refreshing again." ) - if let accountSnapshot { - await localActivityCollector?.selectPartition( - historyPartition.id - ) + if localAnalyticsVisible, + includeLocalActivity || localAnalyticsNeedsLoad, + let accountSnapshot { await refreshLocalActivity( for: accountSnapshot, observedAt: Date(), @@ -329,6 +341,46 @@ final class UsageMonitor: ObservableObject { } } + func setLocalAnalyticsVisible(_ isVisible: Bool) async { + if isVisible { + if !localAnalyticsVisible { + localAnalyticsVisible = true + localAnalyticsNeedsLoad = true + } + guard localAnalyticsNeedsLoad else { return } + } else { + guard localAnalyticsVisible || localAnalyticsNeedsLoad else { + return + } + localAnalyticsVisible = false + localAnalyticsNeedsLoad = false + cancelLocalImport() + localActivityCollection = .unavailable( + "Codex local records are unavailable" + ) + await localActivityCollector?.releaseCachedFacts() + _ = await recalculate() + return + } + while isRefreshing { + guard !Task.isCancelled else { return } + try? await Task.sleep(for: .milliseconds(50)) + } + guard localAnalyticsVisible, + localAnalyticsNeedsLoad, + let accountSnapshot else { + return + } + let identityVerified = sourceState == .available + && historyAccountIdentity != nil + await refreshLocalActivity( + for: accountSnapshot, + observedAt: identityVerified ? accountSnapshot.fetchedAt : Date(), + identityVerified: identityVerified + ) + _ = await recalculate() + } + func updateSafetyBuffer(_ value: Double) { let value = SafetyBufferPolicy.normalized(value) defaults.set(value, forKey: Self.safetyBufferKey) @@ -425,7 +477,6 @@ final class UsageMonitor: ObservableObject { planType: String? = nil, observedAt: Date ) async { - cancelLocalImport() let partition: AccountHistoryPartition let authState: String let previousAuthState = defaults.string(forKey: Self.historyAuthStateKey) @@ -473,8 +524,8 @@ final class UsageMonitor: ObservableObject { if let planType { defaults.set(planType, forKey: Self.historyPlanTypeKey) } - await localActivityCollector?.selectPartition(partition.id) guard partition != historyPartition else { return } + cancelLocalImport() historyPartition = partition historyConnectionActive = false if let data = try? JSONEncoder().encode(partition) { @@ -906,6 +957,8 @@ final class UsageMonitor: ObservableObject { identityVerified: Bool = true ) async { cancelLocalImport() + let generation = localImportGeneration + localAnalyticsNeedsLoad = false guard let interval = UsageIntelligenceEngine.tokenActivityInterval( account: snapshot, samples: historyMatchesCurrentSnapshot ? samples : [], @@ -922,6 +975,8 @@ final class UsageMonitor: ObservableObject { ) return } + await localActivityCollector.selectPartition(historyPartition.id) + guard generation == localImportGeneration else { return } localActivityCollection = .unavailable( "Codex local records are unavailable" ) @@ -929,16 +984,22 @@ final class UsageMonitor: ObservableObject { interval: interval, observedAt: observedAt ) - if await localActivityCollector.hasPendingHistoryDeletion() { + guard generation == localImportGeneration else { return } + let deletionPending = + await localActivityCollector.hasPendingHistoryDeletion() + guard generation == localImportGeneration else { return } + if deletionPending { historyDeletionStatus = .pendingLocal } localActivityCollection = identityVerified ? collection : collection.loweringCoverage( "Codex account identity could not be checked" - ) - guard await localActivityCollector.hasPendingImport() else { return } - let generation = localImportGeneration + ) + let importPending = await localActivityCollector.hasPendingImport() + guard generation == localImportGeneration, importPending else { + return + } localImportTask = Task(priority: .background) { [weak self] in try? await Task.sleep(for: .milliseconds(500)) await self?.continueLocalActivityImport( diff --git a/Tests/CodexLimitsTests/AnalyticsWorkspaceTests.swift b/Tests/CodexLimitsTests/AnalyticsWorkspaceTests.swift index 1b25518..c5129e4 100644 --- a/Tests/CodexLimitsTests/AnalyticsWorkspaceTests.swift +++ b/Tests/CodexLimitsTests/AnalyticsWorkspaceTests.swift @@ -149,6 +149,17 @@ final class AnalyticsWorkspaceTests: XCTestCase { XCTAssertFalse(AnalyticsGraph.tokenActivity.usesAccountScope) } + func testOnlyUsageRemainingSkipsLocalActivity() { + var state = AnalyticsExplorationState.initial + + XCTAssertFalse(state.usesLocalActivity) + state.graph = .tokenActivity + XCTAssertTrue(state.usesLocalActivity) + state.graph = .usageRemaining + state.section = .facts + XCTAssertTrue(state.usesLocalActivity) + } + func testPresetRangeEndsAtLatestObservedTimeAndIsClampedToWindow() { let end = Date(timeIntervalSince1970: 10 * 86_400) let bounds = DateInterval( diff --git a/Tests/CodexLimitsTests/LocalActivityCollectorTests.swift b/Tests/CodexLimitsTests/LocalActivityCollectorTests.swift index 6591098..be5d2db 100644 --- a/Tests/CodexLimitsTests/LocalActivityCollectorTests.swift +++ b/Tests/CodexLimitsTests/LocalActivityCollectorTests.swift @@ -47,6 +47,87 @@ final class LocalActivityCollectorTests: XCTestCase { XCTAssertEqual(second.observation.coverage, .high) } + func testReleasedFactsRestoreFromThePersistedCache() async throws { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: fixture.root.appendingPathComponent( + "collector-state", + isDirectory: true + ) + ) + await collector.selectPartition("stable-account") + let interval = try fixture.interval() + let first = await collector.refresh(interval: interval) + + await collector.releaseCachedFacts() + let restored = await collector.refresh(interval: interval) + + XCTAssertEqual(restored.facts, first.facts) + } + + func testReleasedFactsStayReleasedWhenARefreshWasSuspended() async throws { + let fixture = try CollectorFixture() + let file = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let versionDelay = SecondInstalledVersionDelay() + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: fixture.root.appendingPathComponent( + "collector-state", + isDirectory: true + ), + installedCLIVersion: { + await versionDelay.response() + } + ) + await collector.selectPartition("stable-account") + let interval = try fixture.interval() + let first = await collector.refresh(interval: interval) + try fixture.append( + fixture.tokens(total: 800, ordinal: 3, minute: 3), + to: file + ) + let suspended = Task { + await collector.refresh(interval: interval) + } + await versionDelay.waitUntilSecondRequest() + + await collector.releaseCachedFacts() + await versionDelay.releaseSecondRequest() + _ = await suspended.value + let restored = await collector.refresh( + interval: interval, + refreshMetadata: false + ) + + XCTAssertEqual( + first.facts.filter { $0.key == .token }.compactMap(\.numericDelta), + [500] + ) + XCTAssertEqual( + restored.facts.filter { $0.key == .token } + .compactMap(\.numericDelta), + [500, 200] + ) + } + func testMissingTrackedFileKeepsFactsAndNamesTheSourceGap() async throws { let fixture = try CollectorFixture() let file = try fixture.rollout( @@ -2402,6 +2483,32 @@ private actor CancellableProjectionDelay { } } +private actor SecondInstalledVersionDelay { + private var requestCount = 0 + private var secondRequestStarted = false + private var secondRequestContinuation: CheckedContinuation? + + func response() async -> String? { + requestCount += 1 + guard requestCount == 2 else { return "0.145.0" } + secondRequestStarted = true + return await withCheckedContinuation { continuation in + secondRequestContinuation = continuation + } + } + + func waitUntilSecondRequest() async { + while !secondRequestStarted { + await Task.yield() + } + } + + func releaseSecondRequest() { + secondRequestContinuation?.resume(returning: "0.145.0") + secondRequestContinuation = nil + } +} + private actor SupersededProjectionDelay { private var listStarted = false private var listContinuation: CheckedContinuation? diff --git a/Tests/CodexLimitsTests/UsageMonitorHistoryTests.swift b/Tests/CodexLimitsTests/UsageMonitorHistoryTests.swift index c11db76..fb2156c 100644 --- a/Tests/CodexLimitsTests/UsageMonitorHistoryTests.swift +++ b/Tests/CodexLimitsTests/UsageMonitorHistoryTests.swift @@ -802,6 +802,10 @@ final class UsageMonitorHistoryTests: XCTestCase { ) await firstMonitor.refresh() await firstMonitor.refresh() + XCTAssertNil(firstMonitor.readerSnapshot.localTokenActivity.tokens) + + await firstMonitor.setLocalAnalyticsVisible(true) + XCTAssertEqual( firstMonitor.readerSnapshot.localTokenActivity.tokens, 400 @@ -818,6 +822,9 @@ final class UsageMonitorHistoryTests: XCTestCase { fetchUsage: { throw CodexClientError.invalidResponse } ) await restarted.refresh() + XCTAssertNil(restarted.readerSnapshot.localTokenActivity.tokens) + + await restarted.setLocalAnalyticsVisible(true) XCTAssertEqual(restarted.readerSnapshot.localTokenActivity.tokens, 400) XCTAssertEqual(restarted.readerSnapshot.localTokenActivity.coverage, .low) @@ -827,7 +834,125 @@ final class UsageMonitorHistoryTests: XCTestCase { ) } - func testMonitorFinishesABoundedLocalImportWithoutAnotherAccountRead() + func testLocalCollectorReadsOnlyForVisibleAnalyticsAndManualRefresh() + async throws + { + let suiteName = "UsageMonitorHistoryTests-\(UUID().uuidString)" + let defaults = try XCTUnwrap(UserDefaults(suiteName: suiteName)) + defer { defaults.removePersistentDomain(forName: suiteName) } + let root = temporaryDirectory() + let localRoot = root.appendingPathComponent( + "rollouts", + isDirectory: true + ) + try FileManager.default.createDirectory( + at: localRoot, + withIntermediateDirectories: true + ) + let requests = RequestCounter() + let fetches = DelayedFetchSource( + makeFetchResult( + identity: "user@example.com", + fetchedAt: Date(timeIntervalSince1970: 1_700_000), + remaining: 90 + ) + ) + let monitor = UsageMonitor( + defaults: defaults, + historyDirectory: root.appendingPathComponent("history"), + startsAutomatically: false, + localActivityCollector: LocalActivityCollector( + rootDirectory: localRoot, + stateDirectory: root.appendingPathComponent("local-state"), + projectionSource: ReadOnlyThreadProjectionSource { _ in + await requests.record() + return Data( + #"{"result":{"data":[],"nextCursor":null}}"#.utf8 + ) + } + ), + fetchUsage: { try await fetches.next() } + ) + + await monitor.start() + var requestCount = await requests.count + XCTAssertEqual(requestCount, 0) + + let automatic = Task { @MainActor in + await monitor.automaticRefresh() + } + try await Task.sleep(for: .milliseconds(10)) + await monitor.setLocalAnalyticsVisible(true) + await automatic.value + requestCount = await requests.count + XCTAssertEqual(requestCount, 1) + await monitor.setLocalAnalyticsVisible(true) + requestCount = await requests.count + XCTAssertEqual(requestCount, 1) + + await monitor.automaticRefresh() + requestCount = await requests.count + XCTAssertEqual(requestCount, 1) + await monitor.refresh() + requestCount = await requests.count + XCTAssertEqual(requestCount, 2) + + await monitor.setLocalAnalyticsVisible(false) + await monitor.refresh() + requestCount = await requests.count + XCTAssertEqual(requestCount, 2) + } + + func testHidingAnalyticsDiscardsAnInFlightLocalRead() async throws { + let suiteName = "UsageMonitorHistoryTests-\(UUID().uuidString)" + let defaults = try XCTUnwrap(UserDefaults(suiteName: suiteName)) + defer { defaults.removePersistentDomain(forName: suiteName) } + let root = temporaryDirectory() + let localRoot = root.appendingPathComponent( + "rollouts", + isDirectory: true + ) + try FileManager.default.createDirectory( + at: localRoot, + withIntermediateDirectories: true + ) + let gate = ProjectionGate() + let fetchResult = makeFetchResult( + identity: "user@example.com", + fetchedAt: Date(timeIntervalSince1970: 1_700_000), + remaining: 90 + ) + let monitor = UsageMonitor( + defaults: defaults, + historyDirectory: root.appendingPathComponent("history"), + startsAutomatically: false, + localActivityCollector: LocalActivityCollector( + rootDirectory: localRoot, + stateDirectory: root.appendingPathComponent("local-state"), + projectionSource: ReadOnlyThreadProjectionSource { _ in + await gate.response() + } + ), + fetchUsage: { fetchResult } + ) + await monitor.refresh() + + let load = Task { @MainActor in + await monitor.setLocalAnalyticsVisible(true) + } + await gate.waitUntilStarted() + let hide = Task { @MainActor in + await monitor.setLocalAnalyticsVisible(false) + } + try await Task.sleep(for: .milliseconds(10)) + await gate.release() + await load.value + await hide.value + + XCTAssertNil(monitor.readerSnapshot.localTokenActivity.tokens) + } + + func testMonitorFinishesABoundedLocalImportAcrossAutomaticRefresh() async throws { let suiteName = "UsageMonitorHistoryTests-\(UUID().uuidString)" @@ -876,6 +1001,11 @@ final class UsageMonitorHistoryTests: XCTestCase { identity: "user@example.com", fetchedAt: fetchedAt, remaining: 80 + ), + makeFetchResult( + identity: "user@example.com", + fetchedAt: fetchedAt.addingTimeInterval(60), + remaining: 79 ) ]) let monitor = UsageMonitor( @@ -886,9 +1016,11 @@ final class UsageMonitorHistoryTests: XCTestCase { fetchUsage: { try await fetches.next() } ) + await monitor.setLocalAnalyticsVisible(true) await monitor.refresh() try Data(rollout.utf8).write(to: rolloutURL) await monitor.refresh() + await monitor.automaticRefresh() for _ in 0 ..< 100 { if await collector.hasPendingImport() == false, monitor.readerSnapshot.localTokenActivity.tokens == 1_009_900 { @@ -904,7 +1036,7 @@ final class UsageMonitorHistoryTests: XCTestCase { 1_009_900 ) let accountReadCount = await fetches.callCount - XCTAssertEqual(accountReadCount, 2) + XCTAssertEqual(accountReadCount, 3) } func testMissingWeeklyRefreshClearsWeeklyOutputsAndKeepsOtherLimits() async throws { @@ -1861,6 +1993,43 @@ private actor FetchSequence { var callCount: Int { calls } } +private actor RequestCounter { + private var requests = 0 + + func record() { + requests += 1 + } + + var count: Int { requests } +} + +private actor ProjectionGate { + private var started = false + private var continuation: CheckedContinuation? + + func response() async -> Data { + started = true + return await withCheckedContinuation { continuation in + self.continuation = continuation + } + } + + func waitUntilStarted() async { + while !started { + await Task.yield() + } + } + + func release() { + continuation?.resume( + returning: Data( + #"{"result":{"data":[],"nextCursor":null}}"#.utf8 + ) + ) + continuation = nil + } +} + private actor DelayedFetchSource { private let result: CodexFetchResult private var calls = 0 From a4c5c188bf07578f072eed6f3b1202ad127ac5fe Mon Sep 17 00:00:00 2001 From: thrr87 <193831865+thrr87@users.noreply.github.com> Date: Thu, 30 Jul 2026 16:35:19 +0200 Subject: [PATCH 5/5] perf: store token activity in SQLite --- Sources/CodexLimits/AnalyticsWorkspace.swift | 16 +- .../CodexLimits/LocalActivityCollector.swift | 825 ++++++++- .../CodexLimits/LocalActivityNormalizer.swift | 92 +- Sources/CodexLimits/LocalActivityStore.swift | 1523 +++++++++++++++++ Sources/CodexLimits/MenuContentView.swift | 38 +- Sources/CodexLimits/RolloutTailSource.swift | 136 +- .../CodexLimits/UsageIntelligenceEngine.swift | 15 +- Sources/CodexLimits/UsageMonitor.swift | 266 ++- .../AnalyticsWorkspaceTests.swift | 59 +- .../LocalActivityCollectorTests.swift | 1113 ++++++++++++ .../LocalActivityNormalizerTests.swift | 1 + .../LocalActivityPerformanceTests.swift | 262 +++ .../LocalActivityStoreTests.swift | 904 ++++++++++ .../LocalTokenActivityTests.swift | 55 + .../RolloutTailSourceTests.swift | 26 + .../UsageMonitorHistoryTests.swift | 239 ++- 16 files changed, 5407 insertions(+), 163 deletions(-) create mode 100644 Sources/CodexLimits/LocalActivityStore.swift create mode 100644 Tests/CodexLimitsTests/LocalActivityStoreTests.swift diff --git a/Sources/CodexLimits/AnalyticsWorkspace.swift b/Sources/CodexLimits/AnalyticsWorkspace.swift index 629b580..7720dac 100644 --- a/Sources/CodexLimits/AnalyticsWorkspace.swift +++ b/Sources/CodexLimits/AnalyticsWorkspace.swift @@ -18,6 +18,11 @@ enum AnalyticsGraph: String, CaseIterable, Codable, Identifiable, Sendable { var id: String { rawValue } + static let coreCases: [AnalyticsGraph] = [ + .usageRemaining, + .tokenActivity + ] + var usesAccountScope: Bool { self == .usageRemaining || self == .usagePerToken } @@ -103,7 +108,12 @@ struct AnalyticsExplorationState: Codable, Equatable, Sendable { ) var usesLocalActivity: Bool { - section != .graphs || graph != .usageRemaining + usesStoredTokenActivity + } + + var usesStoredTokenActivity: Bool { + section == .graphs + && graph == .tokenActivity } } @@ -137,6 +147,10 @@ final class AnalyticsWorkspaceStore: ObservableObject { ) else { return .initial } + guard restored.section == .graphs, + AnalyticsGraph.coreCases.contains(restored.graph) else { + return .initial + } return restored } diff --git a/Sources/CodexLimits/LocalActivityCollector.swift b/Sources/CodexLimits/LocalActivityCollector.swift index 09aabd3..cf26d9c 100644 --- a/Sources/CodexLimits/LocalActivityCollector.swift +++ b/Sources/CodexLimits/LocalActivityCollector.swift @@ -48,11 +48,22 @@ struct LocalActivityCollection: Equatable, Sendable { } } +struct LocalStoredTokenActivityCollection: Equatable, Sendable { + let snapshot: LocalTokenActivitySnapshot + let filterOptions: UsageReceiptFilterOptions + let observation: LocalActivityObservation + let bytesRead: UInt64 + let importPending: Bool +} + actor LocalActivityCollector { private static let maximumMetadataBytes = 1_048_576 private static let maximumRefreshLines = 10_000 private static let maximumRefreshBytes: UInt64 = 8 * 1_024 * 1_024 + private static let maximumStoredRefreshLines = 10_000 private static let maximumRolloutRecordBytes = 7 * 1_024 * 1_024 + private static let maximumStoredRolloutRecordBytes = 256 * 1_024 + private static let maximumStoredRefreshBytes: UInt64 = 64 * 1_024 * 1_024 private struct ObservationSignature: Equatable { let sourceVersion: String? @@ -107,6 +118,41 @@ actor LocalActivityCollector { var facts: [LocalActivityFact] } + private struct StoredTokenImportCheckpoint: Codable { + var retainedTokenAfterCutoff: Bool + } + + private struct StoredTokenCacheKey: Equatable { + let interval: DateInterval + let filters: WorkspaceFilters + } + + private struct StoredTokenCache { + let key: StoredTokenCacheKey + let snapshot: LocalTokenActivitySnapshot + let filterOptions: UsageReceiptFilterOptions + } + + private struct StoredFilterOptionsCache { + let interval: DateInterval + let options: UsageReceiptFilterOptions + } + + private struct StoredTokenFactProjection: Decodable { + struct Source: Decodable { + let sourceGeneration: UInt64 + } + + let key: LocalActivityFactKey + let availability: LocalActivityAvailability + let numericDelta: Int64? + let reason: String? + let eventID: String? + let eventTimestamp: String? + let source: Source + let context: LocalActivityContext? + } + private struct DeletionMarker: Codable { let version: Int let cutoff: Date @@ -155,6 +201,9 @@ actor LocalActivityCollector { private var refreshGeneration: UInt64 = 0 private var didReadInstalledCLIVersion = false private var cachedInstalledCLIVersion: String? + private var storedTokenStore: LocalActivityStore? + private var storedTokenCache: StoredTokenCache? + private var storedFilterOptionsCache: StoredFilterOptionsCache? init( rootDirectory: URL = FileManager.default.homeDirectoryForCurrentUser @@ -194,6 +243,7 @@ actor LocalActivityCollector { guard partitionID != id else { return } stateGeneration = nextRevision(after: stateGeneration) persist() + closeStoredTokenStore() partitionID = id files.removeAll() restoredFilesByFingerprint.removeAll() @@ -808,6 +858,709 @@ actor LocalActivityCollector { ) } + func refreshStoredTokenActivity( + interval: DateInterval, + filters: WorkspaceFilters = .all, + observedAt: Date = Date() + ) async -> LocalStoredTokenActivityCollection { + autoreleasepool { + refreshStoredTokenActivitySync( + interval: interval, + filters: filters, + observedAt: observedAt + ) + } + } + + private func refreshStoredTokenActivitySync( + interval: DateInterval, + filters: WorkspaceFilters, + observedAt: Date + ) -> LocalStoredTokenActivityCollection { + guard !historyDeletionPending else { + return unavailableStoredTokenActivity( + deletionMarkerInvalid + ? "Saved deletion state could not be read" + : "Analytics history deletion is pending", + interval: interval + ) + } + guard interval.start < interval.end else { + return unavailableStoredTokenActivity( + "Weekly token interval is unavailable", + interval: interval + ) + } + do { + let store = try openStoredTokenStore() + var result = try importStoredTokenFacts( + into: store, + interval: interval + ) + if result.bytesRead > 0, !result.importPending { + result.importPending = true + result.gapReason = result.gapReason + ?? "Local task import is still in progress" + } else if !result.importPending { + result.merge( + try refreshStoredTokenRolloutTail( + in: interval, + observedAt: observedAt, + store: store + ) + ) + } + guard result.hasSources else { + return unavailableStoredTokenActivity( + "No saved local activity is available", + interval: interval + ) + } + if result.importPending { + let reason = "Local task import is still in progress" + let version = result.sourceVersions.sorted().first ?? "unknown" + return LocalStoredTokenActivityCollection( + snapshot: .unavailable(reason, interval: interval), + filterOptions: UsageReceiptFilterOptions( + projects: [], + taskTrees: [], + models: [], + reasoningLevels: [] + ), + observation: .gap( + sourceVersion: version, + observedAt: observedAt, + reason: reason + ), + bytesRead: result.bytesRead, + importPending: true + ) + } + if !result.importPending, result.gapReason == nil { + if result.sourceVersions.isEmpty { + result.gapReason = "Codex CLI version is unavailable" + } else if result.sourceVersions != Set(["0.145.0"]) { + result.gapReason = + "This Codex CLI version has not been checked" + } else { + result.gapReason = "Local task discovery is incomplete" + } + } + let version = result.sourceVersions.sorted().first ?? "unknown" + let observation: LocalActivityObservation + if let reason = result.gapReason { + observation = .gap( + sourceVersion: version, + observedAt: observedAt, + reason: reason + ) + } else { + observation = .continuous( + sourceVersion: version, + observedAt: observedAt + ) + } + let cacheKey = StoredTokenCacheKey( + interval: interval, + filters: filters + ) + if result.bytesRead == 0, + let cached = storedTokenCache, + cached.key == cacheKey { + return LocalStoredTokenActivityCollection( + snapshot: cached.snapshot.updating( + interval: interval, + observation: observation + ), + filterOptions: cached.filterOptions, + observation: observation, + bytesRead: 0, + importPending: false + ) + } + let snapshot = try store.tokenActivity( + in: interval, + filters: filters, + observation: observation + ) + let filterOptions: UsageReceiptFilterOptions + if result.bytesRead == 0, + let cached = storedFilterOptionsCache, + cached.interval.start <= interval.start, + cached.interval.end >= interval.end { + filterOptions = cached.options + } else { + filterOptions = try store.filterOptions(in: interval) + storedFilterOptionsCache = StoredFilterOptionsCache( + interval: interval, + options: filterOptions + ) + } + storedTokenCache = StoredTokenCache( + key: cacheKey, + snapshot: snapshot, + filterOptions: filterOptions + ) + return LocalStoredTokenActivityCollection( + snapshot: snapshot, + filterOptions: filterOptions, + observation: observation, + bytesRead: result.bytesRead, + importPending: result.importPending + ) + } catch is CancellationError { + return unavailableStoredTokenActivity( + "Local activity read was cancelled", + interval: interval + ) + } catch { + return unavailableStoredTokenActivity( + "Saved local activity could not be read", + interval: interval + ) + } + } + + private struct StoredTokenImportResult { + var bytesRead: UInt64 = 0 + var importPending = false + var hasSources = false + var sourceVersions = Set() + var gapReason: String? + + mutating func merge(_ other: StoredTokenImportResult) { + bytesRead += other.bytesRead + importPending = importPending || other.importPending + hasSources = hasSources || other.hasSources + sourceVersions.formUnion(other.sourceVersions) + if let reason = other.gapReason { + recordGap(reason) + } + } + + mutating func recordGap(_ reason: String) { + guard Self.gapPriority(reason) > Self.gapPriority(gapReason) + else { + return + } + gapReason = reason + } + + private static func gapPriority(_ reason: String?) -> Int { + switch reason { + case "Local task record continuity changed": + 3 + case "Some local diagnostic records could not be read": + 2 + case nil: + 0 + default: + 1 + } + } + } + + private func importStoredTokenFacts( + into store: LocalActivityStore, + interval: DateInterval + ) throws -> StoredTokenImportResult { + guard let directory = stateURL, + let entries = try? FileManager.default.contentsOfDirectory( + at: directory, + includingPropertiesForKeys: nil, + options: [.skipsHiddenFiles] + ) else { + return StoredTokenImportResult() + } + var result = StoredTokenImportResult() + var remainingLines = Self.maximumStoredRefreshLines + var remainingBytes = Self.maximumStoredRefreshBytes + for entry in entries.sorted(by: { $0.path < $1.path }) + where entry.pathExtension == "json" { + guard let data = Self.readMetadata(at: entry), + let persisted = try? JSONDecoder().decode( + PersistedFile.self, + from: data + ), + [4, 5, 6, 7].contains(persisted.version) else { + result.recordGap("Saved local activity could not be read") + continue + } + if let activityStart = persisted.activityStart, + let activityEnd = persisted.activityEnd, + activityStart >= interval.end || activityEnd < interval.start { + continue + } + if let discontinuityAt = persisted.discontinuityAt, + discontinuityAt >= interval.start, + discontinuityAt <= interval.end { + result.recordGap("Local task record continuity changed") + } else if persisted.hasMalformedRecords == true, + persisted.activityStart.map({ $0 < interval.end }) + ?? false, + persisted.activityEnd.map({ $0 >= interval.start }) + ?? false { + result.recordGap( + "Some local diagnostic records could not be read" + ) + } + if persisted.normalization.sourceVersion != "unknown" { + result.sourceVersions.insert( + persisted.normalization.sourceVersion + ) + } + let fingerprint = entry.deletingPathExtension().lastPathComponent + let factsFile = factsURL( + forFingerprint: fingerprint, + in: directory + ) + guard var source = try storedSource( + key: fingerprint, + persisted: persisted, + factsFile: factsFile + ) else { + result.recordGap("Saved local activity could not be read") + continue + } + result.hasSources = true + if try store.hasActiveSource(matching: source) { + continue + } + guard remainingLines > 0, remainingBytes > 0 else { + result.importPending = true + break + } + let replacement: Int64 + if let resumed = try store.resumableReplacement(matching: source) { + replacement = resumed.storeGeneration + source = resumed.source + } else { + replacement = try store.beginReplacement(for: source) + } + let storedProjections = ( + [persisted.projection].compactMap { $0 } + + (persisted.ancestorProjections ?? []) + ).map(\.withoutRolloutFileURL) + do { + let step = try readStoredTokenFacts( + from: factsFile, + source: source, + maximumLines: remainingLines, + maximumBytes: remainingBytes + ) + source = step.source + try store.appendReplacementTokenFacts( + step.facts, + to: source, + storeGeneration: replacement, + projections: storedProjections + ) + remainingLines = max( + remainingLines - step.linesRead, + 0 + ) + remainingBytes = step.bytesRead >= remainingBytes + ? 0 + : remainingBytes - step.bytesRead + let total = result.bytesRead.addingReportingOverflow( + step.bytesRead + ) + result.bytesRead = total.overflow ? .max : total.partialValue + if step.complete { + try store.activateReplacement( + sourceKey: source.key, + storeGeneration: replacement + ) + } else { + result.importPending = true + } + } catch is CancellationError { + throw CancellationError() + } catch { + try? store.rollbackReplacement( + sourceKey: source.key, + storeGeneration: replacement + ) + result.recordGap("Saved local activity could not be read") + } + } + return result + } + + private func refreshStoredTokenRolloutTail( + in interval: DateInterval, + observedAt: Date, + store: LocalActivityStore + ) throws -> StoredTokenImportResult { + guard FileManager.default.fileExists(atPath: rootDirectory.path) else { + return StoredTokenImportResult( + gapReason: "Codex local records are unavailable" + ) + } + let discoveryInterval = DateInterval( + start: interval.start.addingTimeInterval(-86_400), + end: interval.end + ) + let paths = Set( + storedRolloutFiles(in: discoveryInterval).compactMap(fileKey) + ) + guard !paths.isEmpty else { + return StoredTokenImportResult() + } + var result = StoredTokenImportResult(hasSources: true) + var remainingLines = Self.maximumStoredRefreshLines + var remainingBytes = Self.maximumStoredRefreshBytes + + for path in paths.sorted(by: >) { + try Task.checkCancellation() + guard remainingLines > 0, remainingBytes > 0 else { + result.importPending = true + result.recordGap("Local task import is still in progress") + break + } + let file = fileURL(path) + let sourceKey = stateFileName(for: path) + let previousSource = try store.activeSource(key: sourceKey) + let previousCursor: RolloutCursor? + let previousNormalization: LocalActivityNormalizationState? + if let previousSource { + guard let cursorData = previousSource.cursorJSON, + let normalizationData = + previousSource.normalizationJSON, + let cursor = try? JSONDecoder().decode( + RolloutCursor.self, + from: cursorData + ), + let normalization = try? JSONDecoder().decode( + LocalActivityNormalizationState.self, + from: normalizationData + ) else { + result.recordGap("Saved local activity could not be read") + continue + } + previousCursor = cursor + previousNormalization = normalization + } else { + previousCursor = nil + previousNormalization = nil + } + if previousSource?.discontinuityAt != nil { + if let version = previousNormalization?.sourceVersion, + version != "unknown" { + result.sourceVersions.insert(version) + } + result.recordGap("Local task record continuity changed") + continue + } + let batch: RolloutTailBatch + do { + batch = try tail.read( + fileURL: file, + cursor: previousCursor, + observedAt: observedAt, + recordScope: .tokenActivity, + maximumLines: remainingLines, + maximumBytes: remainingBytes, + maximumRecordBytes: + Self.maximumStoredRolloutRecordBytes + ) + } catch is CancellationError { + throw CancellationError() + } catch { + result.recordGap("Local task records are missing") + continue + } + remainingLines = max( + remainingLines - batch.processedLineCount, + 0 + ) + remainingBytes = batch.bytesRead >= remainingBytes + ? 0 + : remainingBytes - batch.bytesRead + result.bytesRead += batch.bytesRead + if batch.malformedRecordCount > 0 { + result.recordGap( + "Some local diagnostic records could not be read" + ) + } + if previousSource != nil, + batch.requiresRebuild || batch.continuityChanged { + if var previousSource { + previousSource.discontinuityAt = observedAt + try store.appendTokenFacts([], to: previousSource) + } + result.recordGap("Local task record continuity changed") + continue + } + if batch.hasMoreRecords { + result.importPending = true + result.recordGap("Local task import is still in progress") + } + if batch.records.isEmpty, + batch.cursor == previousCursor, + batch.malformedRecordCount == 0 { + if let version = previousNormalization?.sourceVersion, + version != "unknown" { + result.sourceVersions.insert(version) + } + continue + } + captureRolloutProjections( + batch.records, + sourceGeneration: batch.cursor.sourceGeneration, + observedAt: observedAt + ) + let normalized = normalizer.normalize( + records: recordsAfterHistoryCutoff(batch.records), + sourceGeneration: batch.cursor.sourceGeneration, + observedAt: observedAt, + previousState: previousNormalization + ) + if normalized.state.sourceVersion != "unknown" { + result.sourceVersions.insert(normalized.state.sourceVersion) + } + let newFacts = factsAfterHistoryCutoff(normalized.facts).filter { + $0.key == .token + && $0.availability == .available + && $0.eventID != nil + } + var source = previousSource ?? LocalActivityStore.Source( + key: sourceKey, + sourceGeneration: batch.cursor.sourceGeneration + ) + source.cursorJSON = try JSONEncoder().encode(batch.cursor) + source.normalizationJSON = try JSONEncoder().encode( + normalized.state + ) + source.hasMalformedRecords = + source.hasMalformedRecords || batch.malformedRecordCount > 0 + if let bounds = tokenActivityBounds(newFacts) { + source.activityStart = source.activityStart.map { + min($0, bounds.start) + } ?? bounds.start + source.activityEnd = source.activityEnd.map { + max($0, bounds.end) + } ?? bounds.end + } + let storedProjections = + normalized.state.context?.taskID.map { + projectionChainIDs(for: $0).compactMap { + projections[$0]?.withoutRolloutFileURL + } + } ?? [] + try store.append( + newFacts, + to: source, + projections: storedProjections + ) + } + return result + } + + private func storedSource( + key: String, + persisted: PersistedFile, + factsFile: URL + ) throws -> LocalActivityStore.Source? { + guard let attributes = try? FileManager.default.attributesOfItem( + atPath: factsFile.path + ), + let device = (attributes[.systemNumber] as? NSNumber)?.uint64Value, + let inode = (attributes[.systemFileNumber] as? NSNumber)?.uint64Value, + let size = (attributes[.size] as? NSNumber)?.uint64Value, + let modificationDate = attributes[.modificationDate] as? Date else { + return nil + } + var source = LocalActivityStore.Source( + key: key, + sourceGeneration: persisted.cursor.sourceGeneration + ) + source.cursorJSON = try JSONEncoder().encode(persisted.cursor) + source.normalizationJSON = try JSONEncoder().encode( + persisted.normalization + ) + source.activityStart = persisted.activityStart + source.activityEnd = persisted.activityEnd + source.discontinuityAt = persisted.discontinuityAt + source.hasMalformedRecords = persisted.hasMalformedRecords ?? false + source.legacyDevice = device + source.legacyInode = inode + source.legacyOffset = 0 + source.legacySize = size + source.legacyModificationDate = modificationDate + return source + } + + private func readStoredTokenFacts( + from file: URL, + source: LocalActivityStore.Source, + maximumLines: Int, + maximumBytes: UInt64 + ) throws -> ( + facts: [LocalActivityStore.TokenFact], + source: LocalActivityStore.Source, + linesRead: Int, + bytesRead: UInt64, + complete: Bool + ) { + guard let fileSize = source.legacySize, + let handle = try? FileHandle(forReadingFrom: file) else { + throw CocoaError(.fileReadUnknown) + } + defer { try? handle.close() } + let decoder = JSONDecoder() + let checkpoint: StoredTokenImportCheckpoint + if let data = source.pendingFactJSON { + checkpoint = try decoder.decode( + StoredTokenImportCheckpoint.self, + from: data + ) + } else { + checkpoint = StoredTokenImportCheckpoint( + retainedTokenAfterCutoff: false + ) + } + var retainedTokenAfterCutoff = + checkpoint.retainedTokenAfterCutoff + var facts: [LocalActivityStore.TokenFact] = [] + var valid = true + let read = try BoundedJSONLReader.read( + handle: handle, + from: source.legacyOffset ?? 0, + maximumLines: maximumLines, + maximumBytes: maximumBytes, + maximumRecordBytes: Self.maximumMetadataBytes, + discardsPartialRecordAtByteLimit: false + ) { line, _ in + guard let decoded = autoreleasepool(invoking: { + try? decoder.decode( + StoredTokenFactProjection.self, + from: line + ) + }) else { + valid = false + return false + } + guard decoded.key == .token, + decoded.availability == .available, + let eventID = decoded.eventID, + !eventID.isEmpty, + let timestamp = decoded.eventTimestamp, + let occurredAt = parseTimestamp(timestamp), + historyCutoff.map({ occurredAt >= $0 }) ?? true else { + return true + } + var numericDelta = decoded.numericDelta + var reason = decoded.reason + if historyCutoff != nil, + !retainedTokenAfterCutoff { + numericDelta = nil + reason = reason ?? "segment-baseline" + retainedTokenAfterCutoff = true + } + facts.append( + LocalActivityStore.TokenFact( + eventID: eventID, + occurredAt: occurredAt, + numericDelta: numericDelta, + reason: reason, + sourceGeneration: decoded.source.sourceGeneration, + taskID: decoded.context?.taskID, + effectiveModel: decoded.context?.effectiveModel, + reasoning: decoded.context?.reasoning + ) + ) + return true + } + guard valid, read.oversizedRecordCount == 0 else { + throw CocoaError(.fileReadCorruptFile) + } + let complete = read.completeByteOffset == fileSize + guard complete || read.stoppedEarly else { + throw CocoaError(.fileReadCorruptFile) + } + var updated = source + if complete { + updated.legacyOffset = fileSize + updated.pendingFactJSON = nil + } else { + updated.legacyOffset = read.resumeByteOffset + updated.pendingFactJSON = try JSONEncoder().encode( + StoredTokenImportCheckpoint( + retainedTokenAfterCutoff: retainedTokenAfterCutoff + ) + ) + } + return ( + facts, + updated, + read.processedLineCount, + read.bytesRead, + complete + ) + } + + private func openStoredTokenStore() throws -> LocalActivityStore { + guard partitionID != nil, let directory = stateURL else { + throw CocoaError(.fileNoSuchFile) + } + if let storedTokenStore { + return storedTokenStore + } + let file = directory.appendingPathComponent( + "local-activity-v1.sqlite3" + ) + do { + let store = try LocalActivityStore( + fileURL: file, + historyCutoff: historyCutoff + ) + storedTokenStore = store + return store + } catch let error as LocalActivityStore.StoreError + where error == .historyCutoffMismatch || error == .schemaMismatch { + for candidate in [ + file, + URL(fileURLWithPath: file.path + "-wal"), + URL(fileURLWithPath: file.path + "-shm"), + URL(fileURLWithPath: file.path + "-journal") + ] where FileManager.default.fileExists(atPath: candidate.path) { + try FileManager.default.removeItem(at: candidate) + } + let store = try LocalActivityStore( + fileURL: file, + historyCutoff: historyCutoff + ) + storedTokenStore = store + return store + } + } + + private func closeStoredTokenStore() { + storedTokenStore?.close() + storedTokenStore = nil + storedTokenCache = nil + storedFilterOptionsCache = nil + } + + private func unavailableStoredTokenActivity( + _ reason: String, + interval: DateInterval + ) -> LocalStoredTokenActivityCollection { + LocalStoredTokenActivityCollection( + snapshot: .unavailable(reason, interval: interval), + filterOptions: UsageReceiptFilterOptions( + projects: [], + taskTrees: [], + models: [], + reasoningLevels: [] + ), + observation: .unavailable(reason), + bytesRead: 0, + importPending: false + ) + } + private func advanceContentRevision() { contentRevision = nextRevision(after: contentRevision) } @@ -842,6 +1595,7 @@ actor LocalActivityCollector { for: stateDirectory ) } + closeStoredTokenStore() stateGeneration = nextRevision(after: stateGeneration) files.removeAll() restoredFilesByFingerprint.removeAll() @@ -887,6 +1641,7 @@ actor LocalActivityCollector { ) historyCutoff = cutoff deletionMarkerInvalid = false + closeStoredTokenStore() stateGeneration = nextRevision(after: stateGeneration) files.removeAll() restoredFilesByFingerprint.removeAll() @@ -915,6 +1670,7 @@ actor LocalActivityCollector { } func rebuildHistory() throws { + closeStoredTokenStore() if let stateDirectory { if FileManager.default.fileExists(atPath: stateDirectory.path) { try FileManager.default.removeItem(at: stateDirectory) @@ -941,14 +1697,17 @@ actor LocalActivityCollector { deletionMarkerInvalid = false } - private func rolloutFiles(in interval: DateInterval) -> Set { + private func rolloutFiles( + in interval: DateInterval, + maximumDays: Int = 14 + ) -> Set { var calendar = Calendar(identifier: .gregorian) calendar.timeZone = TimeZone(secondsFromGMT: 0)! var date = calendar.startOfDay(for: interval.start) let end = calendar.startOfDay(for: interval.end) var result = Set() - for _ in 0 ..< 14 where date <= end { + for _ in 0 ..< max(maximumDays, 1) where date <= end { let components = calendar.dateComponents( [.year, .month, .day], from: date @@ -983,6 +1742,32 @@ actor LocalActivityCollector { return result } + private func storedRolloutFiles( + in interval: DateInterval + ) -> Set { + var calendar = Calendar(identifier: .gregorian) + calendar.timeZone = TimeZone(secondsFromGMT: 0)! + let end = calendar.startOfDay(for: interval.end) + guard let start = calendar.date( + byAdding: .day, + value: -85, + to: end + ) else { + return [] + } + return rolloutFiles( + in: DateInterval(start: start, end: end), + maximumDays: 86 + ).filter { + guard let modified = try? $0.resourceValues( + forKeys: [.contentModificationDateKey] + ).contentModificationDate else { + return true + } + return modified >= interval.start + } + } + private func taskIDs(in states: [FileState]) -> Set { Set(states.compactMap(taskID(in:))) } @@ -1199,6 +1984,40 @@ actor LocalActivityCollector { }.last } + private func captureRolloutProjections( + _ records: [RolloutRecord], + sourceGeneration: UInt64, + observedAt: Date + ) { + for record in records + where record.type == "session_meta" { + guard let taskID = record.threadID, + projections[taskID] == nil else { + continue + } + let label = record.cwd.flatMap { + let value = URL(fileURLWithPath: $0).lastPathComponent + return value.isEmpty ? nil : value + } + projections[taskID] = ThreadProjection( + taskID: taskID, + parentTaskID: record.parentThreadID, + projectLabel: label, + rolloutFileURL: nil, + createdAt: record.timestamp.flatMap(parseTimestamp), + updatedAt: nil, + source: LocalActivitySourceMetadata( + source: .rolloutJSONL, + sourceVersion: record.cliVersion ?? "unknown", + schemaVersion: "rollout-jsonl-v1", + sourceGeneration: sourceGeneration, + historyMode: record.historyMode, + observedAt: observedAt + ) + ) + } + } + private func activeProjections( in interval: DateInterval ) -> [ThreadProjection] { @@ -1249,6 +2068,8 @@ actor LocalActivityCollector { lastPublishedProjectionIdentities = nil lastObservationSignature = nil importContinuationPending = false + storedTokenCache = nil + storedFilterOptionsCache = nil } private func restartPartialFactRestores() { diff --git a/Sources/CodexLimits/LocalActivityNormalizer.swift b/Sources/CodexLimits/LocalActivityNormalizer.swift index df4aa2f..fb8c842 100644 --- a/Sources/CodexLimits/LocalActivityNormalizer.swift +++ b/Sources/CodexLimits/LocalActivityNormalizer.swift @@ -98,6 +98,9 @@ struct LocalActivityNormalizationState: Codable, Equatable, Sendable { var lastTokenUsage: LocalTokenUsage? = nil var tokenSegment: UInt64 var context: LocalActivityContext? = nil + var primaryTaskID: String? = nil + var awaitsPrimaryTaskStart: Bool? = nil + var usesNextContextTokenUsage: Bool? = nil } struct LocalActivityNormalizationResult: Equatable, Sendable { @@ -143,6 +146,9 @@ struct LocalActivityNormalizer { state.tokenSegment += 1 } state.context = nil + state.primaryTaskID = nil + state.awaitsPrimaryTaskStart = nil + state.usesNextContextTokenUsage = nil } state.sourceGeneration = sourceGeneration if let observedVersion { @@ -168,6 +174,31 @@ struct LocalActivityNormalizer { var currentModelContextWindow = state.context?.modelContextWindow for record in records { + if state.primaryTaskID == nil, + record.type == "session_meta", + let taskID = record.threadID { + state.primaryTaskID = taskID + state.awaitsPrimaryTaskStart = + record.parentThreadID == nil ? nil : true + } + if state.awaitsPrimaryTaskStart == true { + guard isPrimaryTaskStart( + record, + taskID: state.primaryTaskID + ) else { + continue + } + state.awaitsPrimaryTaskStart = nil + state.lastTotalTokens = nil + state.lastTokenUsage = nil + state.usesNextContextTokenUsage = true + currentTaskID = state.primaryTaskID + currentTurnID = nil + currentAgent = nil + currentModel = nil + currentReasoning = nil + currentModelContextWindow = nil + } if let taskID = record.threadID { currentTaskID = taskID } @@ -286,13 +317,15 @@ struct LocalActivityNormalizer { } if let tokenUsage = record.tokenUsage { let totalTokens = tokenUsage.totalTokens - let delta = state.lastTotalTokens.flatMap { previous in - totalTokens >= previous - ? difference(totalTokens, previous) + let delta = state.lastTotalTokens.flatMap { + totalTokens >= $0 ? difference(totalTokens, $0) : nil + } ?? ( + state.usesNextContextTokenUsage == true + ? record.contextTokenUsage?.totalTokens : nil - } + ) let reason: String? - if state.lastTotalTokens == nil { + if state.lastTotalTokens == nil, delta == nil { reason = sourceChanged ? "source-discontinuity" : "segment-baseline" } else if delta == nil { reason = "cumulative-counter-decreased" @@ -310,7 +343,12 @@ struct LocalActivityNormalizer { numericDelta: delta, tokenSegment: state.tokenSegment, reason: reason, - eventID: record.eventID, + eventID: tokenEventID( + record: record, + context: context, + totalTokens: totalTokens, + tokenSegment: state.tokenSegment + ), eventTimestamp: record.timestamp, source: source, context: context, @@ -322,6 +360,7 @@ struct LocalActivityNormalizer { ) state.lastTotalTokens = totalTokens state.lastTokenUsage = tokenUsage + state.usesNextContextTokenUsage = nil } if record.tokenUsage == nil, let contextTokenUsage = record.contextTokenUsage { @@ -417,6 +456,47 @@ struct LocalActivityNormalizer { ) } + private func tokenEventID( + record: RolloutRecord, + context: LocalActivityContext, + totalTokens: Int64, + tokenSegment: UInt64 + ) -> String { + guard let taskID = context.taskID, + let turnID = context.turnID else { + return record.eventID + } + return [ + "token", + taskID, + turnID, + String(tokenSegment), + String(totalTokens) + ].joined(separator: "|") + } + + private func isPrimaryTaskStart( + _ record: RolloutRecord, + taskID: String? + ) -> Bool { + guard record.eventType == "task_started", + let taskTimestamp = taskID.flatMap(uuidV7Timestamp), + let turnTimestamp = record.turnID.flatMap(uuidV7Timestamp) else { + return false + } + return turnTimestamp >= taskTimestamp + } + + private func uuidV7Timestamp(_ value: String) -> UInt64? { + let compact = value.replacingOccurrences(of: "-", with: "") + guard compact.count == 32, + compact[compact.index(compact.startIndex, offsetBy: 12)] == "7" + else { + return nil + } + return UInt64(compact.prefix(12), radix: 16) + } + private func tokenDelta( from previous: LocalTokenUsage, to current: LocalTokenUsage diff --git a/Sources/CodexLimits/LocalActivityStore.swift b/Sources/CodexLimits/LocalActivityStore.swift new file mode 100644 index 0000000..876bb42 --- /dev/null +++ b/Sources/CodexLimits/LocalActivityStore.swift @@ -0,0 +1,1523 @@ +import Foundation +import SQLite3 + +final class LocalActivityStore { + struct TokenFact: Equatable, Sendable { + let eventID: String + let occurredAt: Date + let numericDelta: Int64? + let reason: String? + let sourceGeneration: UInt64 + let taskID: String? + let effectiveModel: String? + let reasoning: String? + + init( + eventID: String, + occurredAt: Date, + numericDelta: Int64?, + reason: String?, + sourceGeneration: UInt64, + taskID: String? = nil, + effectiveModel: String? = nil, + reasoning: String? = nil + ) { + self.eventID = eventID + self.occurredAt = occurredAt + self.numericDelta = numericDelta + self.reason = reason + self.sourceGeneration = sourceGeneration + self.taskID = taskID + self.effectiveModel = effectiveModel + self.reasoning = reasoning + } + } + + struct Source: Equatable, Sendable { + let key: String + let sourceGeneration: UInt64 + var cursorJSON: Data? = nil + var normalizationJSON: Data? = nil + var activityStart: Date? = nil + var activityEnd: Date? = nil + var discontinuityAt: Date? = nil + var hasMalformedRecords = false + var legacyDevice: UInt64? = nil + var legacyInode: UInt64? = nil + var legacyOffset: UInt64? = nil + var legacySize: UInt64? = nil + var legacyModificationDate: Date? = nil + var pendingFactJSON: Data? = nil + } + + enum StoreError: Swift.Error, Equatable { + case closed + case invalidQuery + case invalidSource + case sourceGenerationChanged + case historyCutoffMismatch + case schemaMismatch + case tokenTotalOverflow + case sqlite(String) + } + + private static let schemaVersion = 2 + private static let taskTreeCTE = """ + WITH RECURSIVE task_tree( + root_task_id, task_id, project_label, depth + ) AS ( + SELECT task_id, task_id, project_label, 0 + FROM task_projection + WHERE parent_task_id IS NULL + UNION ALL + SELECT tree.root_task_id, + child.task_id, + tree.project_label, + tree.depth + 1 + FROM task_projection AS child + JOIN task_tree AS tree + ON child.parent_task_id = tree.task_id + WHERE tree.depth < 64 + ) + """ + private static let transient = unsafeBitCast( + -1, + to: sqlite3_destructor_type.self + ) + + private var database: OpaquePointer? + var testFactProgress: ((Int) -> Void)? + + init(fileURL: URL, historyCutoff: Date? = nil) throws { + try FileManager.default.createDirectory( + at: fileURL.deletingLastPathComponent(), + withIntermediateDirectories: true + ) + var opened: OpaquePointer? + let result = sqlite3_open_v2( + fileURL.path, + &opened, + SQLITE_OPEN_CREATE | SQLITE_OPEN_READWRITE | SQLITE_OPEN_FULLMUTEX, + nil + ) + guard result == SQLITE_OK, let opened else { + let message = opened.map { + String(cString: sqlite3_errmsg($0)) + } ?? "Could not open local activity store" + if let opened { + sqlite3_close_v2(opened) + } + throw StoreError.sqlite(message) + } + database = opened + do { + try execute("PRAGMA journal_mode = WAL") + try execute("PRAGMA synchronous = NORMAL") + try execute("PRAGMA foreign_keys = ON") + try execute("PRAGMA cache_size = -4096") + try execute("PRAGMA mmap_size = 0") + try execute("PRAGMA temp_store = FILE") + try execute("PRAGMA busy_timeout = 1000") + try execute("PRAGMA journal_size_limit = 8388608") + try installSchema(historyCutoff: historyCutoff) + } catch { + close() + throw error + } + } + + deinit { + close() + } + + func close() { + guard let database else { return } + sqlite3_close_v2(database) + self.database = nil + } + + func append( + _ facts: [LocalActivityFact], + to source: Source, + projections: [ThreadProjection] = [] + ) throws { + guard !source.key.isEmpty else { + throw StoreError.invalidSource + } + try transaction { + let storeGeneration = try activeStoreGeneration(for: source) + ?? createActiveSource(source) + try upsert(projections) + try appendFacts( + tokenFacts(from: facts, source: source), + to: source, + storeGeneration: storeGeneration + ) + } + } + + func appendTokenFacts( + _ facts: [TokenFact], + to source: Source, + projections: [ThreadProjection] = [] + ) throws { + guard !source.key.isEmpty else { + throw StoreError.invalidSource + } + try transaction { + let storeGeneration = try activeStoreGeneration(for: source) + ?? createActiveSource(source) + try upsert(projections) + try appendFacts( + facts, + to: source, + storeGeneration: storeGeneration + ) + } + } + + @discardableResult + func beginReplacement(for source: Source) throws -> Int64 { + guard !source.key.isEmpty else { + throw StoreError.invalidSource + } + return try transaction { + if let resumed = try importingReplacement(matching: source) { + return resumed.storeGeneration + } + let cleanup = try prepare( + """ + DELETE FROM source + WHERE source_key = ?1 AND is_active = 0 + """ + ) + defer { sqlite3_finalize(cleanup) } + try bind(source.key, at: 1, to: cleanup) + try stepDone(cleanup) + let storeGeneration = try nextStoreGeneration(sourceKey: source.key) + try insertSource( + source, + storeGeneration: storeGeneration, + isActive: false + ) + try update(source, storeGeneration: storeGeneration) + return storeGeneration + } + } + + func resumableReplacement( + matching source: Source + ) throws -> (storeGeneration: Int64, source: Source)? { + try importingReplacement(matching: source) + } + + func hasActiveSource(matching source: Source) throws -> Bool { + let statement = try prepare( + """ + SELECT EXISTS( + SELECT 1 + FROM source + WHERE source_key = ?1 + AND is_active = 1 + AND source_generation = ?2 + AND legacy_device IS ?3 + AND legacy_inode IS ?4 + AND legacy_size IS ?5 + AND legacy_mtime IS ?6 + AND pending_fact_json IS NULL + AND ( + ( + legacy_device IS NULL + AND legacy_inode IS NULL + AND legacy_offset IS NULL + AND legacy_size IS NULL + AND legacy_mtime IS NULL + ) + OR ( + legacy_device IS NOT NULL + AND legacy_inode IS NOT NULL + AND legacy_offset = legacy_size + AND legacy_size IS NOT NULL + AND legacy_mtime IS NOT NULL + ) + ) + ) + """ + ) + defer { sqlite3_finalize(statement) } + try bind(source.key, at: 1, to: statement) + try bind(String(source.sourceGeneration), at: 2, to: statement) + try bind(source.legacyDevice.map(String.init), at: 3, to: statement) + try bind(source.legacyInode.map(String.init), at: 4, to: statement) + try bind(source.legacySize.map(String.init), at: 5, to: statement) + try bind(source.legacyModificationDate, at: 6, to: statement) + guard sqlite3_step(statement) == SQLITE_ROW else { + throw StoreError.sqlite(errorMessage) + } + return sqlite3_column_int(statement, 0) != 0 + } + + func hasActiveSource(key: String) throws -> Bool { + guard !key.isEmpty else { throw StoreError.invalidSource } + let statement = try prepare( + """ + SELECT EXISTS( + SELECT 1 FROM source + WHERE source_key = ?1 AND is_active = 1 + ) + """ + ) + defer { sqlite3_finalize(statement) } + try bind(key, at: 1, to: statement) + guard sqlite3_step(statement) == SQLITE_ROW else { + throw StoreError.sqlite(errorMessage) + } + return sqlite3_column_int(statement, 0) != 0 + } + + func activeSource(key: String) throws -> Source? { + guard !key.isEmpty else { throw StoreError.invalidSource } + let statement = try prepare( + """ + SELECT source_generation, cursor_json, normalization_json, + activity_start, activity_end, discontinuity_at, malformed, + legacy_device, legacy_inode, legacy_offset, legacy_size, + legacy_mtime, pending_fact_json + FROM source + WHERE source_key = ?1 AND is_active = 1 + """ + ) + defer { sqlite3_finalize(statement) } + try bind(key, at: 1, to: statement) + switch sqlite3_step(statement) { + case SQLITE_DONE: + return nil + case SQLITE_ROW: + guard let sourceGeneration = UInt64( + columnString(statement, at: 0) ?? "" + ) else { + throw StoreError.sqlite("Stored source generation is invalid") + } + var source = Source( + key: key, + sourceGeneration: sourceGeneration + ) + source.cursorJSON = columnData(statement, at: 1) + source.normalizationJSON = columnData(statement, at: 2) + source.activityStart = columnDate(statement, at: 3) + source.activityEnd = columnDate(statement, at: 4) + source.discontinuityAt = columnDate(statement, at: 5) + source.hasMalformedRecords = + sqlite3_column_int(statement, 6) != 0 + source.legacyDevice = columnUInt64(statement, at: 7) + source.legacyInode = columnUInt64(statement, at: 8) + source.legacyOffset = columnUInt64(statement, at: 9) + source.legacySize = columnUInt64(statement, at: 10) + source.legacyModificationDate = columnDate(statement, at: 11) + source.pendingFactJSON = columnData(statement, at: 12) + return source + default: + throw StoreError.sqlite(errorMessage) + } + } + + func appendReplacement( + _ facts: [LocalActivityFact], + to source: Source, + storeGeneration: Int64, + projections: [ThreadProjection] = [] + ) throws { + try transaction { + guard try isImporting( + sourceKey: source.key, + storeGeneration: storeGeneration, + sourceGeneration: source.sourceGeneration + ) else { + throw StoreError.invalidSource + } + try upsert(projections) + try appendFacts( + tokenFacts(from: facts, source: source), + to: source, + storeGeneration: storeGeneration + ) + } + } + + func appendReplacementTokenFacts( + _ facts: [TokenFact], + to source: Source, + storeGeneration: Int64, + projections: [ThreadProjection] = [] + ) throws { + try transaction { + guard try isImporting( + sourceKey: source.key, + storeGeneration: storeGeneration, + sourceGeneration: source.sourceGeneration + ) else { + throw StoreError.invalidSource + } + try upsert(projections) + try appendFacts( + facts, + to: source, + storeGeneration: storeGeneration + ) + } + } + + func activateReplacement( + sourceKey: String, + storeGeneration: Int64 + ) throws { + try transaction { + guard try isCompleteImport( + sourceKey: sourceKey, + storeGeneration: storeGeneration + ) else { + throw StoreError.invalidSource + } + let retire = try prepare( + """ + UPDATE source + SET is_active = 0 + WHERE source_key = ?1 AND is_active = 1 + """ + ) + defer { sqlite3_finalize(retire) } + try bind(sourceKey, at: 1, to: retire) + try stepDone(retire) + let activate = try prepare( + """ + UPDATE source + SET is_active = 1 + WHERE source_key = ?1 AND store_generation = ?2 + """ + ) + defer { sqlite3_finalize(activate) } + try bind(sourceKey, at: 1, to: activate) + try bind(storeGeneration, at: 2, to: activate) + try stepDone(activate) + try rebuildCanonicalTokens(affectedBy: sourceKey) + let cleanup = try prepare( + """ + DELETE FROM source + WHERE source_key = ?1 AND store_generation != ?2 + """ + ) + defer { sqlite3_finalize(cleanup) } + try bind(sourceKey, at: 1, to: cleanup) + try bind(storeGeneration, at: 2, to: cleanup) + try stepDone(cleanup) + } + } + + func rollbackReplacement( + sourceKey: String, + storeGeneration: Int64 + ) throws { + try transaction { + let statement = try prepare( + """ + DELETE FROM source + WHERE source_key = ?1 + AND store_generation = ?2 + AND is_active = 0 + """ + ) + defer { sqlite3_finalize(statement) } + try bind(sourceKey, at: 1, to: statement) + try bind(storeGeneration, at: 2, to: statement) + try stepDone(statement) + } + } + + func tokenActivity( + in interval: DateInterval, + filters: WorkspaceFilters = .all, + observation: LocalActivityObservation, + maximumPoints: Int = 1_000 + ) throws -> LocalTokenActivitySnapshot { + guard let database else { throw StoreError.closed } + try Task.checkCancellation() + sqlite3_progress_handler( + database, + 1_000, + { _ in Task.isCancelled ? 1 : 0 }, + nil + ) + defer { sqlite3_progress_handler(database, 0, nil, nil) } + guard maximumPoints > 0, interval.start < interval.end else { + throw StoreError.invalidQuery + } + if case let .unavailable(reason) = observation { + return .unavailable(reason, interval: interval) + } + let bucketWidth = interval.duration / Double(maximumPoints) + let pointsSQL = Self.taskTreeCTE + """ + , filtered AS ( + SELECT event.* + FROM canonical_token AS canonical + JOIN token_fact AS event + ON event.source_id = canonical.source_id + AND event.event_id = canonical.event_id + LEFT JOIN task_tree + ON task_tree.task_id = event.task_id + WHERE event.occurred_at >= ?1 + AND event.occurred_at < ?2 + AND event.numeric_delta >= 0 + AND (?5 IS NULL OR task_tree.project_label = ?5) + AND (?6 IS NULL OR task_tree.root_task_id = ?6) + AND (?7 IS NULL OR event.effective_model = ?7) + AND (?8 IS NULL OR event.reasoning = ?8) + ), + bucketed AS ( + SELECT MIN( + CAST((occurred_at - ?1) / ?3 AS INTEGER), + ?4 + ) AS bucket, + MAX(occurred_at) AS occurred_at, + SUM(numeric_delta) AS delta + FROM filtered + GROUP BY bucket + ), + running AS ( + SELECT occurred_at, + SUM(delta) OVER (ORDER BY bucket) AS tokens + FROM bucketed + ) + SELECT occurred_at, tokens + FROM running + ORDER BY occurred_at + """ + let statement = try prepare(pointsSQL) + defer { sqlite3_finalize(statement) } + try bind(interval.start.timeIntervalSince1970, at: 1, to: statement) + try bind(interval.end.timeIntervalSince1970, at: 2, to: statement) + try bind(bucketWidth, at: 3, to: statement) + try bind(Int64(maximumPoints - 1), at: 4, to: statement) + try bind(filters.projectID, at: 5, to: statement) + try bind(filters.taskTreeID, at: 6, to: statement) + try bind(filters.model, at: 7, to: statement) + try bind(filters.reasoning, at: 8, to: statement) + var points: [LocalTokenActivityPoint] = [] + while true { + switch sqlite3_step(statement) { + case SQLITE_ROW: + points.append( + LocalTokenActivityPoint( + date: Date( + timeIntervalSince1970: + sqlite3_column_double(statement, 0) + ), + tokens: sqlite3_column_int64(statement, 1) + ) + ) + case SQLITE_DONE: + break + case SQLITE_INTERRUPT: + throw CancellationError() + default: + let message = errorMessage + if message.localizedCaseInsensitiveContains("overflow") { + throw StoreError.tokenTotalOverflow + } + throw StoreError.sqlite(message) + } + if sqlite3_data_count(statement) == 0 { + break + } + } + let hasUnboundedCounter = try scalarBool( + Self.taskTreeCTE + """ + SELECT EXISTS( + SELECT 1 + FROM canonical_token AS canonical + JOIN token_fact AS event + ON event.source_id = canonical.source_id + AND event.event_id = canonical.event_id + LEFT JOIN task_tree + ON task_tree.task_id = event.task_id + WHERE event.occurred_at >= ?1 + AND event.occurred_at < ?2 + AND event.break_kind != 0 + AND ( + event.numeric_delta IS NULL + OR event.numeric_delta < 0 + ) + AND (?3 IS NULL OR task_tree.project_label = ?3) + AND (?4 IS NULL OR task_tree.root_task_id = ?4) + AND (?5 IS NULL OR event.effective_model = ?5) + AND (?6 IS NULL OR event.reasoning = ?6) + ) + """, + interval: interval, + filters: filters + ) + let coverage: CoverageLevel + let reason: String? + switch observation { + case .continuous: + if hasUnboundedCounter { + coverage = .low + reason = "Local token activity starts from an unbounded counter" + } else if points.isEmpty { + coverage = .notApplicable + reason = "No local token activity was observed" + } else { + coverage = .high + reason = "Only local activity on this Mac is observed" + } + case let .gap(_, _, message): + coverage = .low + reason = message + case let .unavailable(message): + coverage = .unavailable + reason = message + } + let source: (String?, Date?) = switch observation { + case let .continuous(version, observedAt), + let .gap(version, observedAt, _): + (version, observedAt) + case .unavailable: + (nil, nil) + } + try Task.checkCancellation() + return LocalTokenActivitySnapshot( + tokens: points.last?.tokens ?? 0, + interval: interval, + coverage: coverage, + reason: reason, + sourceVersion: source.0, + observedAt: source.1, + points: points, + accountComparison: .unavailable + ) + } + + func filterOptions( + in interval: DateInterval + ) throws -> UsageReceiptFilterOptions { + guard let database else { throw StoreError.closed } + try Task.checkCancellation() + sqlite3_progress_handler( + database, + 1_000, + { _ in Task.isCancelled ? 1 : 0 }, + nil + ) + defer { sqlite3_progress_handler(database, 0, nil, nil) } + guard interval.start < interval.end else { + throw StoreError.invalidQuery + } + let statement = try prepare( + Self.taskTreeCTE + """ + SELECT DISTINCT + task_tree.project_label, + task_tree.root_task_id, + event.effective_model, + event.reasoning + FROM canonical_token AS canonical + JOIN token_fact AS event + ON event.source_id = canonical.source_id + AND event.event_id = canonical.event_id + LEFT JOIN task_tree ON task_tree.task_id = event.task_id + WHERE event.occurred_at >= ?1 + AND event.occurred_at < ?2 + AND event.numeric_delta > 0 + """ + ) + defer { sqlite3_finalize(statement) } + try bind(interval.start, at: 1, to: statement) + try bind(interval.end, at: 2, to: statement) + var projects = Set() + var taskTrees = Set() + var models = Set() + var reasoningLevels = Set() + while true { + switch sqlite3_step(statement) { + case SQLITE_ROW: + if let project = columnString(statement, at: 0) { + projects.insert(project) + } + if let taskTree = columnString(statement, at: 1) { + taskTrees.insert(taskTree) + } + if let model = columnString(statement, at: 2) { + models.insert(model) + } + if let reasoning = columnString(statement, at: 3) { + reasoningLevels.insert(reasoning) + } + case SQLITE_DONE: + return UsageReceiptFilterOptions( + projects: projects.sorted(), + taskTrees: taskTrees.sorted(), + models: models.sorted(), + reasoningLevels: reasoningLevels.sorted() + ) + case SQLITE_INTERRUPT: + throw CancellationError() + default: + throw StoreError.sqlite(errorMessage) + } + } + } + + private func installSchema(historyCutoff: Date?) throws { + let version = try scalarInt("PRAGMA user_version") + guard version == 0 || version == Self.schemaVersion else { + throw StoreError.schemaMismatch + } + if version == Self.schemaVersion { + try validateHistoryCutoff(historyCutoff) + return + } + try execute("PRAGMA auto_vacuum = INCREMENTAL") + try transaction { + try execute( + """ + CREATE TABLE store_meta ( + id INTEGER PRIMARY KEY CHECK (id = 1), + history_cutoff REAL + ) + """ + ) + let metadata = try prepare( + "INSERT INTO store_meta (id, history_cutoff) VALUES (1, ?1)" + ) + defer { sqlite3_finalize(metadata) } + try bind(historyCutoff, at: 1, to: metadata) + try stepDone(metadata) + try execute( + """ + CREATE TABLE source ( + id INTEGER PRIMARY KEY, + source_key TEXT NOT NULL, + store_generation INTEGER NOT NULL, + source_generation TEXT NOT NULL, + is_active INTEGER NOT NULL, + cursor_json BLOB, + normalization_json BLOB, + activity_start REAL, + activity_end REAL, + discontinuity_at REAL, + malformed INTEGER NOT NULL, + legacy_device TEXT, + legacy_inode TEXT, + legacy_offset TEXT, + legacy_size TEXT, + legacy_mtime REAL, + pending_fact_json BLOB, + UNIQUE (source_key, store_generation) + ) + """ + ) + try execute( + """ + CREATE UNIQUE INDEX one_active_source + ON source(source_key) + WHERE is_active = 1 + """ + ) + try execute( + """ + CREATE TABLE token_fact ( + source_id INTEGER NOT NULL, + event_id TEXT NOT NULL, + occurred_at REAL NOT NULL, + numeric_delta INTEGER, + break_kind INTEGER NOT NULL, + task_id TEXT, + effective_model TEXT, + reasoning TEXT, + PRIMARY KEY (source_id, event_id), + FOREIGN KEY (source_id) + REFERENCES source(id) + ON DELETE CASCADE + ) WITHOUT ROWID + """ + ) + try execute( + """ + CREATE TABLE canonical_token ( + event_id TEXT PRIMARY KEY, + source_id INTEGER NOT NULL, + FOREIGN KEY (source_id) + REFERENCES source(id) + ON DELETE CASCADE + ) WITHOUT ROWID + """ + ) + try execute( + """ + CREATE INDEX token_fact_time + ON token_fact(occurred_at) + """ + ) + try execute( + """ + CREATE TABLE task_projection ( + task_id TEXT PRIMARY KEY, + parent_task_id TEXT, + project_label TEXT, + created_at REAL, + updated_at REAL, + source_kind TEXT NOT NULL + ) WITHOUT ROWID + """ + ) + try execute( + """ + CREATE INDEX task_projection_parent + ON task_projection(parent_task_id) + """ + ) + try execute("PRAGMA user_version = \(Self.schemaVersion)") + } + } + + private func validateHistoryCutoff(_ historyCutoff: Date?) throws { + let statement = try prepare( + "SELECT history_cutoff FROM store_meta WHERE id = 1" + ) + defer { sqlite3_finalize(statement) } + guard sqlite3_step(statement) == SQLITE_ROW, + columnDate(statement, at: 0) == historyCutoff else { + throw StoreError.historyCutoffMismatch + } + } + + private func activeStoreGeneration(for source: Source) throws -> Int64? { + let statement = try prepare( + """ + SELECT store_generation, source_generation + FROM source + WHERE source_key = ?1 AND is_active = 1 + """ + ) + defer { sqlite3_finalize(statement) } + try bind(source.key, at: 1, to: statement) + switch sqlite3_step(statement) { + case SQLITE_DONE: + return nil + case SQLITE_ROW: + let sourceGeneration = String( + cString: sqlite3_column_text(statement, 1) + ) + guard sourceGeneration == String(source.sourceGeneration) else { + throw StoreError.sourceGenerationChanged + } + return sqlite3_column_int64(statement, 0) + default: + throw StoreError.sqlite(errorMessage) + } + } + + private func sourceRecord( + sourceKey: String, + storeGeneration: Int64 + ) throws -> (id: Int64, isActive: Bool) { + let statement = try prepare( + """ + SELECT id, is_active + FROM source + WHERE source_key = ?1 AND store_generation = ?2 + """ + ) + defer { sqlite3_finalize(statement) } + try bind(sourceKey, at: 1, to: statement) + try bind(storeGeneration, at: 2, to: statement) + let result = sqlite3_step(statement) + if result == SQLITE_INTERRUPT { + throw CancellationError() + } + guard result == SQLITE_ROW else { + throw StoreError.invalidSource + } + return ( + sqlite3_column_int64(statement, 0), + sqlite3_column_int(statement, 1) != 0 + ) + } + + private func rebuildCanonicalTokens(affectedBy sourceKey: String) throws { + let delete = try prepare( + """ + DELETE FROM canonical_token + WHERE event_id IN ( + SELECT fact.event_id + FROM token_fact AS fact + JOIN source ON source.id = fact.source_id + WHERE source.source_key = ?1 + ) + """ + ) + defer { sqlite3_finalize(delete) } + try bind(sourceKey, at: 1, to: delete) + try stepDone(delete) + + let insert = try prepare( + """ + INSERT INTO canonical_token (event_id, source_id) + SELECT event_id, source_id + FROM ( + SELECT fact.event_id, + fact.source_id, + ROW_NUMBER() OVER ( + PARTITION BY fact.event_id + ORDER BY source.source_key, + source.store_generation + ) AS duplicate_order + FROM token_fact AS fact + JOIN source ON source.id = fact.source_id + WHERE source.is_active = 1 + AND fact.event_id IN ( + SELECT affected.event_id + FROM token_fact AS affected + JOIN source AS changed + ON changed.id = affected.source_id + WHERE changed.source_key = ?1 + ) + ) + WHERE duplicate_order = 1 + """ + ) + defer { sqlite3_finalize(insert) } + try bind(sourceKey, at: 1, to: insert) + try stepDone(insert) + } + + private func createActiveSource(_ source: Source) throws -> Int64 { + let storeGeneration = try nextStoreGeneration(sourceKey: source.key) + try insertSource( + source, + storeGeneration: storeGeneration, + isActive: true + ) + return storeGeneration + } + + private func insertSource( + _ source: Source, + storeGeneration: Int64, + isActive: Bool + ) throws { + let statement = try prepare( + """ + INSERT INTO source ( + source_key, store_generation, source_generation, + is_active, malformed + ) VALUES (?1, ?2, ?3, ?4, 0) + """ + ) + defer { sqlite3_finalize(statement) } + try bind(source.key, at: 1, to: statement) + try bind(storeGeneration, at: 2, to: statement) + try bind(String(source.sourceGeneration), at: 3, to: statement) + try bind(isActive ? 1 : 0, at: 4, to: statement) + try stepDone(statement) + } + + private func isImporting( + sourceKey: String, + storeGeneration: Int64, + sourceGeneration: UInt64? + ) throws -> Bool { + let statement = try prepare( + """ + SELECT source_generation + FROM source + WHERE source_key = ?1 + AND store_generation = ?2 + AND is_active = 0 + """ + ) + defer { sqlite3_finalize(statement) } + try bind(sourceKey, at: 1, to: statement) + try bind(storeGeneration, at: 2, to: statement) + guard sqlite3_step(statement) == SQLITE_ROW else { return false } + guard let sourceGeneration else { return true } + return String(cString: sqlite3_column_text(statement, 0)) + == String(sourceGeneration) + } + + private func isCompleteImport( + sourceKey: String, + storeGeneration: Int64 + ) throws -> Bool { + let statement = try prepare( + """ + SELECT EXISTS( + SELECT 1 + FROM source + WHERE source_key = ?1 + AND store_generation = ?2 + AND is_active = 0 + AND pending_fact_json IS NULL + AND ( + ( + legacy_device IS NULL + AND legacy_inode IS NULL + AND legacy_offset IS NULL + AND legacy_size IS NULL + AND legacy_mtime IS NULL + ) + OR ( + legacy_device IS NOT NULL + AND legacy_inode IS NOT NULL + AND legacy_offset = legacy_size + AND legacy_size IS NOT NULL + AND legacy_mtime IS NOT NULL + ) + ) + ) + """ + ) + defer { sqlite3_finalize(statement) } + try bind(sourceKey, at: 1, to: statement) + try bind(storeGeneration, at: 2, to: statement) + let result = sqlite3_step(statement) + if result == SQLITE_INTERRUPT { + throw CancellationError() + } + guard result == SQLITE_ROW else { + throw StoreError.sqlite(errorMessage) + } + return sqlite3_column_int(statement, 0) != 0 + } + + private func importingReplacement( + matching source: Source + ) throws -> (storeGeneration: Int64, source: Source)? { + let statement = try prepare( + """ + SELECT store_generation, source_generation, + cursor_json, normalization_json, + activity_start, activity_end, discontinuity_at, malformed, + legacy_device, legacy_inode, legacy_offset, legacy_size, + legacy_mtime, pending_fact_json + FROM source + WHERE source_key = ?1 + AND is_active = 0 + AND source_generation = ?2 + AND legacy_device IS ?3 + AND legacy_inode IS ?4 + AND legacy_size IS ?5 + AND legacy_mtime IS ?6 + """ + ) + defer { sqlite3_finalize(statement) } + try bind(source.key, at: 1, to: statement) + try bind(String(source.sourceGeneration), at: 2, to: statement) + try bind(source.legacyDevice.map(String.init), at: 3, to: statement) + try bind(source.legacyInode.map(String.init), at: 4, to: statement) + try bind(source.legacySize.map(String.init), at: 5, to: statement) + try bind(source.legacyModificationDate, at: 6, to: statement) + guard sqlite3_step(statement) == SQLITE_ROW else { return nil } + let storeGeneration = sqlite3_column_int64(statement, 0) + guard let sourceGeneration = UInt64(columnString(statement, at: 1) ?? "") + else { + throw StoreError.sqlite("Stored source generation is invalid") + } + var restored = Source( + key: source.key, + sourceGeneration: sourceGeneration + ) + restored.cursorJSON = columnData(statement, at: 2) + restored.normalizationJSON = columnData(statement, at: 3) + restored.activityStart = columnDate(statement, at: 4) + restored.activityEnd = columnDate(statement, at: 5) + restored.discontinuityAt = columnDate(statement, at: 6) + restored.hasMalformedRecords = sqlite3_column_int(statement, 7) != 0 + restored.legacyDevice = columnUInt64(statement, at: 8) + restored.legacyInode = columnUInt64(statement, at: 9) + restored.legacyOffset = columnUInt64(statement, at: 10) + restored.legacySize = columnUInt64(statement, at: 11) + restored.legacyModificationDate = columnDate(statement, at: 12) + restored.pendingFactJSON = columnData(statement, at: 13) + return (storeGeneration, restored) + } + + private func upsert(_ projections: [ThreadProjection]) throws { + guard !projections.isEmpty else { return } + let statement = try prepare( + """ + INSERT INTO task_projection ( + task_id, parent_task_id, project_label, + created_at, updated_at, source_kind + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6) + ON CONFLICT(task_id) DO UPDATE SET + parent_task_id = excluded.parent_task_id, + project_label = excluded.project_label, + created_at = excluded.created_at, + updated_at = excluded.updated_at, + source_kind = excluded.source_kind + WHERE COALESCE( + excluded.updated_at, + excluded.created_at, + 0 + ) >= COALESCE( + task_projection.updated_at, + task_projection.created_at, + 0 + ) + """ + ) + defer { sqlite3_finalize(statement) } + for projection in projections { + guard !projection.taskID.isEmpty else { continue } + try reset(statement) + try bind(projection.taskID, at: 1, to: statement) + try bind(projection.parentTaskID, at: 2, to: statement) + try bind(projection.projectLabel, at: 3, to: statement) + try bind(projection.createdAt, at: 4, to: statement) + try bind(projection.updatedAt, at: 5, to: statement) + try bind(projection.source.source.rawValue, at: 6, to: statement) + try stepDone(statement) + } + } + + private func appendFacts( + _ facts: [TokenFact], + to source: Source, + storeGeneration: Int64 + ) throws { + try update(source, storeGeneration: storeGeneration) + let sourceRecord = try sourceRecord( + sourceKey: source.key, + storeGeneration: storeGeneration + ) + let factStatement = try prepare( + """ + INSERT OR IGNORE INTO token_fact ( + source_id, event_id, occurred_at, numeric_delta, break_kind, + task_id, effective_model, reasoning + ) VALUES ( + ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8 + ) + """ + ) + defer { sqlite3_finalize(factStatement) } + let canonicalStatement = sourceRecord.isActive + ? try prepare( + """ + INSERT INTO canonical_token (event_id, source_id) + VALUES (?1, ?2) + ON CONFLICT(event_id) DO UPDATE SET + source_id = excluded.source_id + WHERE ( + SELECT candidate.source_key < current.source_key + OR ( + candidate.source_key = current.source_key + AND candidate.store_generation + < current.store_generation + ) + FROM source AS candidate, source AS current + WHERE candidate.id = excluded.source_id + AND current.id = canonical_token.source_id + ) + """ + ) + : nil + defer { sqlite3_finalize(canonicalStatement) } + for (index, fact) in facts.enumerated() { + if index.isMultiple(of: 256) { + try Task.checkCancellation() + testFactProgress?(index) + } + guard fact.sourceGeneration == source.sourceGeneration else { + throw StoreError.sourceGenerationChanged + } + try insertTokenFact( + fact, + sourceID: sourceRecord.id, + statement: factStatement + ) + guard sqlite3_changes(database) > 0, + let canonicalStatement else { + continue + } + try upsertCanonicalToken( + fact, + sourceID: sourceRecord.id, + statement: canonicalStatement + ) + } + } + + private func tokenFacts( + from facts: [LocalActivityFact], + source: Source + ) throws -> [TokenFact] { + let timestampParser = LocalEventTimestampParser() + var tokens: [TokenFact] = [] + tokens.reserveCapacity(facts.count) + for fact in facts { + guard fact.source.sourceGeneration == source.sourceGeneration + else { + throw StoreError.sourceGenerationChanged + } + guard fact.key == .token, + fact.availability == .available, + let eventID = fact.eventID, + !eventID.isEmpty, + let timestamp = fact.eventTimestamp, + let occurredAt = timestampParser.date(from: timestamp) else { + continue + } + tokens.append( + TokenFact( + eventID: eventID, + occurredAt: occurredAt, + numericDelta: fact.numericDelta, + reason: fact.reason, + sourceGeneration: fact.source.sourceGeneration, + taskID: fact.context?.taskID, + effectiveModel: fact.context?.effectiveModel, + reasoning: fact.context?.reasoning + ) + ) + } + return tokens + } + + private func update( + _ source: Source, + storeGeneration: Int64 + ) throws { + let statement = try prepare( + """ + UPDATE source + SET cursor_json = ?3, + normalization_json = ?4, + activity_start = ?5, + activity_end = ?6, + discontinuity_at = ?7, + malformed = ?8, + legacy_device = ?9, + legacy_inode = ?10, + legacy_offset = ?11, + legacy_size = ?12, + legacy_mtime = ?13, + pending_fact_json = ?14 + WHERE source_key = ?1 AND store_generation = ?2 + """ + ) + defer { sqlite3_finalize(statement) } + try bind(source.key, at: 1, to: statement) + try bind(storeGeneration, at: 2, to: statement) + try bind(source.cursorJSON, at: 3, to: statement) + try bind(source.normalizationJSON, at: 4, to: statement) + try bind(source.activityStart, at: 5, to: statement) + try bind(source.activityEnd, at: 6, to: statement) + try bind(source.discontinuityAt, at: 7, to: statement) + try bind(source.hasMalformedRecords ? 1 : 0, at: 8, to: statement) + try bind(source.legacyDevice.map(String.init), at: 9, to: statement) + try bind(source.legacyInode.map(String.init), at: 10, to: statement) + try bind(source.legacyOffset.map(String.init), at: 11, to: statement) + try bind(source.legacySize.map(String.init), at: 12, to: statement) + try bind(source.legacyModificationDate, at: 13, to: statement) + try bind(source.pendingFactJSON, at: 14, to: statement) + try stepDone(statement) + } + + private func nextStoreGeneration(sourceKey: String) throws -> Int64 { + let statement = try prepare( + """ + SELECT COALESCE(MAX(store_generation), 0) + 1 + FROM source WHERE source_key = ?1 + """ + ) + defer { sqlite3_finalize(statement) } + try bind(sourceKey, at: 1, to: statement) + let result = sqlite3_step(statement) + if result == SQLITE_INTERRUPT { + throw CancellationError() + } + guard result == SQLITE_ROW else { + throw StoreError.sqlite(errorMessage) + } + return sqlite3_column_int64(statement, 0) + } + + private func insertTokenFact( + _ fact: TokenFact, + sourceID: Int64, + statement: OpaquePointer + ) throws { + try reset(statement) + try bind(sourceID, at: 1, to: statement) + try bind(fact.eventID, at: 2, to: statement) + try bind(fact.occurredAt, at: 3, to: statement) + try bind(fact.numericDelta, at: 4, to: statement) + try bind(breakKind(fact.reason), at: 5, to: statement) + try bindTokenDetails(fact, to: statement) + try stepDone(statement) + } + + private func upsertCanonicalToken( + _ fact: TokenFact, + sourceID: Int64, + statement: OpaquePointer + ) throws { + try reset(statement) + try bind(fact.eventID, at: 1, to: statement) + try bind(sourceID, at: 2, to: statement) + try stepDone(statement) + } + + private func bindTokenDetails( + _ fact: TokenFact, + to statement: OpaquePointer + ) throws { + try bind(fact.taskID, at: 6, to: statement) + try bind(fact.effectiveModel, at: 7, to: statement) + try bind(fact.reasoning, at: 8, to: statement) + } + + private func breakKind(_ reason: String?) -> Int64 { + switch reason { + case "segment-baseline": + 1 + case "source-discontinuity": + 2 + case "cumulative-counter-decreased": + 3 + default: + 0 + } + } + + private func reset(_ statement: OpaquePointer) throws { + guard sqlite3_reset(statement) == SQLITE_OK, + sqlite3_clear_bindings(statement) == SQLITE_OK else { + throw StoreError.sqlite(errorMessage) + } + } + + private func scalarInt(_ sql: String) throws -> Int { + let statement = try prepare(sql) + defer { sqlite3_finalize(statement) } + guard sqlite3_step(statement) == SQLITE_ROW else { + throw StoreError.sqlite(errorMessage) + } + return Int(sqlite3_column_int64(statement, 0)) + } + + private func columnString( + _ statement: OpaquePointer, + at index: Int32 + ) -> String? { + guard sqlite3_column_type(statement, index) != SQLITE_NULL, + let value = sqlite3_column_text(statement, index) else { + return nil + } + return String(cString: value) + } + + private func columnUInt64( + _ statement: OpaquePointer, + at index: Int32 + ) -> UInt64? { + columnString(statement, at: index).flatMap(UInt64.init) + } + + private func columnDate( + _ statement: OpaquePointer, + at index: Int32 + ) -> Date? { + guard sqlite3_column_type(statement, index) != SQLITE_NULL else { + return nil + } + return Date(timeIntervalSince1970: sqlite3_column_double(statement, index)) + } + + private func columnData( + _ statement: OpaquePointer, + at index: Int32 + ) -> Data? { + guard sqlite3_column_type(statement, index) != SQLITE_NULL else { + return nil + } + let count = Int(sqlite3_column_bytes(statement, index)) + guard count > 0, let bytes = sqlite3_column_blob(statement, index) else { + return Data() + } + return Data(bytes: bytes, count: count) + } + + private func scalarBool( + _ sql: String, + interval: DateInterval, + filters: WorkspaceFilters = .all + ) throws -> Bool { + let statement = try prepare(sql) + defer { sqlite3_finalize(statement) } + try bind(interval.start.timeIntervalSince1970, at: 1, to: statement) + try bind(interval.end.timeIntervalSince1970, at: 2, to: statement) + try bind(filters.projectID, at: 3, to: statement) + try bind(filters.taskTreeID, at: 4, to: statement) + try bind(filters.model, at: 5, to: statement) + try bind(filters.reasoning, at: 6, to: statement) + let result = sqlite3_step(statement) + if result == SQLITE_INTERRUPT { + throw CancellationError() + } + guard result == SQLITE_ROW else { + throw StoreError.sqlite(errorMessage) + } + return sqlite3_column_int(statement, 0) != 0 + } + + private func transaction( + _ body: () throws -> Value + ) throws -> Value { + guard let database else { throw StoreError.closed } + try Task.checkCancellation() + sqlite3_progress_handler( + database, + 1_000, + { _ in Task.isCancelled ? 1 : 0 }, + nil + ) + defer { sqlite3_progress_handler(database, 0, nil, nil) } + try execute("BEGIN IMMEDIATE") + do { + let value = try body() + try Task.checkCancellation() + try execute("COMMIT") + return value + } catch { + try? execute("ROLLBACK") + throw error + } + } + + private func execute(_ sql: String) throws { + guard let database else { throw StoreError.closed } + var message: UnsafeMutablePointer? + let result = sqlite3_exec(database, sql, nil, nil, &message) + guard result == SQLITE_OK else { + let text = message.map { String(cString: $0) } ?? errorMessage + sqlite3_free(message) + throw StoreError.sqlite(text) + } + } + + private func prepare(_ sql: String) throws -> OpaquePointer { + guard let database else { throw StoreError.closed } + var statement: OpaquePointer? + guard sqlite3_prepare_v2(database, sql, -1, &statement, nil) == SQLITE_OK, + let statement else { + throw StoreError.sqlite(errorMessage) + } + return statement + } + + private func stepDone(_ statement: OpaquePointer) throws { + let result = sqlite3_step(statement) + if result == SQLITE_INTERRUPT { + throw CancellationError() + } + guard result == SQLITE_DONE else { + throw StoreError.sqlite(errorMessage) + } + } + + private func bind( + _ value: String?, + at index: Int32, + to statement: OpaquePointer + ) throws { + let result = if let value { + sqlite3_bind_text( + statement, + index, + value, + -1, + Self.transient + ) + } else { + sqlite3_bind_null(statement, index) + } + try checkBinding(result) + } + + private func bind( + _ value: Data?, + at index: Int32, + to statement: OpaquePointer + ) throws { + let result = if let value { + value.withUnsafeBytes { + sqlite3_bind_blob( + statement, + index, + $0.baseAddress, + Int32(value.count), + Self.transient + ) + } + } else { + sqlite3_bind_null(statement, index) + } + try checkBinding(result) + } + + private func bind( + _ value: Date?, + at index: Int32, + to statement: OpaquePointer + ) throws { + if let value { + try bind(value.timeIntervalSince1970, at: index, to: statement) + } else { + try checkBinding(sqlite3_bind_null(statement, index)) + } + } + + private func bind( + _ value: Double, + at index: Int32, + to statement: OpaquePointer + ) throws { + try checkBinding(sqlite3_bind_double(statement, index, value)) + } + + private func bind( + _ value: Int?, + at index: Int32, + to statement: OpaquePointer + ) throws { + if let value { + try bind(Int64(value), at: index, to: statement) + } else { + try checkBinding(sqlite3_bind_null(statement, index)) + } + } + + private func bind( + _ value: Int64?, + at index: Int32, + to statement: OpaquePointer + ) throws { + if let value { + try checkBinding(sqlite3_bind_int64(statement, index, value)) + } else { + try checkBinding(sqlite3_bind_null(statement, index)) + } + } + + private func checkBinding(_ result: Int32) throws { + guard result == SQLITE_OK else { + throw StoreError.sqlite(errorMessage) + } + } + + private var errorMessage: String { + guard let database else { return "Local activity store is closed" } + return String(cString: sqlite3_errmsg(database)) + } +} diff --git a/Sources/CodexLimits/MenuContentView.swift b/Sources/CodexLimits/MenuContentView.swift index b9a1267..f0b4872 100644 --- a/Sources/CodexLimits/MenuContentView.swift +++ b/Sources/CodexLimits/MenuContentView.swift @@ -45,24 +45,6 @@ struct MenuContentView: View { Divider() - Picker( - "View", - selection: Binding( - get: { workspace.state.section }, - set: workspace.selectSection - ) - ) { - ForEach(AnalyticsSection.allCases) { section in - Text(section.rawValue).tag(section) - } - } - .pickerStyle(.segmented) - .labelsHidden() - .padding(.horizontal, 20) - .padding(.vertical, 12) - - Divider() - ScrollView { workspaceContent .padding(20) @@ -77,7 +59,8 @@ struct MenuContentView: View { .frame(width: layout.width, height: layout.height) .task(id: workspace.state.usesLocalActivity) { await monitor.setLocalAnalyticsVisible( - workspace.state.usesLocalActivity + workspace.state.usesLocalActivity, + exploration: workspace.state ) } .onDisappear { @@ -495,7 +478,7 @@ private struct GraphsWorkspace: View { set: store.selectGraph ) ) { - ForEach(AnalyticsGraph.allCases) { graph in + ForEach(AnalyticsGraph.coreCases) { graph in Text(graph.rawValue).tag(graph) } } @@ -1537,13 +1520,7 @@ private struct TokenActivityWorkspace: View { } private var localSlice: LocalTokenActivitySlice { - guard !store.state.filters.isEmpty else { - return reader.localTokenActivity.slice(in: visibleRange) - } - return reader.usageReceipts.localTokenSlice( - in: visibleRange, - filters: store.state.filters - ) + reader.localTokenActivity.slice(in: visibleRange) } private var accountCoversVisibleRange: Bool { @@ -1913,6 +1890,10 @@ private struct TokenActivityWorkspace: View { "Only activity on this Mac is included" case "Saved local activity could not be read": "Saved local activity could not be read" + case "Some local diagnostic records could not be read": + "Some local activity could not be read" + case "Saved deletion state could not be read": + "Saved history settings could not be read" case "Local activity could not be saved": "Local activity could not be saved" case "Local task import is still in progress": @@ -2041,7 +2022,8 @@ private struct WorkspaceFilterMenu: View { } private var options: UsageReceiptFilterOptions { - reader.usageReceipts.filterOptions(in: range) + reader.localFilterOptions + ?? reader.usageReceipts.filterOptions(in: range) } var body: some View { diff --git a/Sources/CodexLimits/RolloutTailSource.swift b/Sources/CodexLimits/RolloutTailSource.swift index 4e23575..d461386 100644 --- a/Sources/CodexLimits/RolloutTailSource.swift +++ b/Sources/CodexLimits/RolloutTailSource.swift @@ -30,6 +30,7 @@ struct BoundedJSONLReader { discardsPartialRecordAtByteLimit: Bool = true, onLine: (Data, UInt64) throws -> Bool ) throws -> BoundedJSONLReadResult { + try Task.checkCancellation() let lineLimit = max(maximumLines, 1) let byteLimit = max(maximumBytes, 1) try handle.seek(toOffset: startOffset) @@ -45,6 +46,7 @@ struct BoundedJSONLReader { var countedCurrentOversizedLine = startsByDiscardingOversizedRecord while true { + try Task.checkCancellation() let (currentOffset, currentOffsetOverflowed) = startOffset.addingReportingOverflow(bytesRead) guard !currentOffsetOverflowed else { @@ -83,7 +85,9 @@ struct BoundedJSONLReader { Int(min(budgetRemaining, UInt64(Self.chunkBytes))) ) guard count > 0, - let chunk = try handle.read(upToCount: count), + let chunk = try autoreleasepool(invoking: { + try handle.read(upToCount: count) + }), !chunk.isEmpty else { break } @@ -139,6 +143,9 @@ struct BoundedJSONLReader { } processedBytes = nextProcessedBytes processedLineCount += 1 + if processedLineCount.isMultiple(of: 256) { + try Task.checkCancellation() + } line.removeAll(keepingCapacity: true) skipsCurrentLine = false countedCurrentOversizedLine = false @@ -396,6 +403,7 @@ struct RolloutRecord: Equatable, Sendable { let eventType: String? let threadID: String? let parentThreadID: String? + let cwd: String? let cliVersion: String? let historyMode: String? let agentRole: String? @@ -427,12 +435,20 @@ struct RolloutTailBatch: Equatable, Sendable { let processedLineCount: Int } +enum RolloutTailRecordScope: Sendable { + case all + case tokenActivity +} + struct IncrementalRolloutTailSource { private static let fingerprintOffsetBasis: UInt64 = 14_695_981_039_346_656_037 private static let fingerprintPrime: UInt64 = 1_099_511_628_211 private static let payloadMarker = Data(#","payload":"#.utf8) + private static let typeValueMarker = Data(#""type":""#.utf8) private static let compactedTypeMarker = Data(#""type":"compacted""#.utf8) private static let emptyPayloadSuffix = Data(#","payload":{}}"#.utf8) + private static let maximumHeaderBytes = 1_024 + private static let maximumPayloadHeaderBytes = 256 // ponytail: bound one batch; stream normalization if one-pass import matters. static let maximumLinesPerBatch = 100_000 static let maximumBytesPerBatch: UInt64 = 64 * 1_024 * 1_024 @@ -441,6 +457,7 @@ struct IncrementalRolloutTailSource { fileURL: URL, cursor: RolloutCursor?, observedAt: Date, + recordScope: RolloutTailRecordScope = .all, maximumLines: Int = Self.maximumLinesPerBatch, maximumBytes: UInt64 = Self.maximumBytesPerBatch, maximumRecordBytes: Int = BoundedJSONLReader.maximumRecordBytes @@ -489,7 +506,8 @@ struct IncrementalRolloutTailSource { handle: handle, fileSize: fileSize, cursor: cursor, - maximumBytes: byteBudget + maximumBytes: byteBudget, + recordScope: recordScope ) } if continuationVerification.matches == nil, @@ -563,14 +581,25 @@ struct IncrementalRolloutTailSource { && cursor?.discardingOversizedRecord == true, discardsPartialRecordAtByteLimit: false ) { line, absoluteLineOffset in - guard let wire = decodeWire(line, decoder: decoder) else { + if recordScope == .tokenActivity, + autoreleasepool(invoking: { + tokenActivityHeaderDisposition(line) + }) == .skip { + unsupportedRecordCount += 1 + return true + } + guard let wire = autoreleasepool(invoking: { + decodeWire(line, decoder: decoder) + }) else { unsupportedRecordCount += 1 malformedRecordCount += 1 return true } let type = wire.type let payload = wire.payload - if type == "session_meta", let observedThreadID = payload.id { + if type == "session_meta", + threadID == nil, + let observedThreadID = payload.id { threadID = observedThreadID } let nonContentIdentity = replayIdentity( @@ -618,11 +647,14 @@ struct IncrementalRolloutTailSource { ].joined(separator: "|") } guard seenEventKeys.insert(eventKey).inserted else { return true } - guard isSupported( - type: type, - eventType: eventType, - toolClass: payload.item?.type - ) else { + let supported = recordScope == .tokenActivity + ? isTokenActivityRecord(type: type, eventType: eventType) + : isSupported( + type: type, + eventType: eventType, + toolClass: payload.item?.type + ) + guard supported else { unsupportedRecordCount += 1 return true } @@ -654,6 +686,7 @@ struct IncrementalRolloutTailSource { eventType: eventType, threadID: threadID, parentThreadID: payload.parentThreadId, + cwd: payload.cwd, cliVersion: payload.cliVersion, historyMode: payload.historyMode, agentRole: payload.agentRole, @@ -729,7 +762,8 @@ struct IncrementalRolloutTailSource { handle: FileHandle, fileSize: UInt64, cursor: RolloutCursor?, - maximumBytes: UInt64 + maximumBytes: UInt64, + recordScope: RolloutTailRecordScope ) throws -> (matches: Bool?, bytesRead: UInt64) { guard let cursor, @@ -750,10 +784,15 @@ struct IncrementalRolloutTailSource { from: 0, through: cursor.byteOffset ) { line, absoluteLineOffset in + if recordScope == .tokenActivity, + tokenActivityHeaderDisposition(line) == .skip { + return true + } guard let wire = decodeWire(line, decoder: decoder) else { return true } if wire.type == "session_meta", + threadID == nil, let observedThreadID = wire.payload.id { threadID = observedThreadID } @@ -872,6 +911,7 @@ struct IncrementalRolloutTailSource { components.append(payload.type ?? "none") components.append(payload.id ?? "none") components.append(payload.parentThreadId ?? "none") + components.append(payload.cwd ?? "none") components.append(payload.cliVersion ?? "none") components.append(payload.historyMode ?? "none") components.append(payload.agentRole ?? "none") @@ -1005,6 +1045,76 @@ struct IncrementalRolloutTailSource { return normalizedToolClass(toolClass) != nil } + private func isTokenActivityRecord( + type: String, + eventType: String? + ) -> Bool { + type == "session_meta" + || type == "turn_context" + || ( + type == "event_msg" + && ["task_started", "token_count"].contains(eventType) + ) + } + + private func tokenActivityHeaderDisposition( + _ data: Data + ) -> TokenActivityHeaderDisposition { + let headerEnd = data.index( + data.startIndex, + offsetBy: min(data.count, Self.maximumHeaderBytes) + ) + guard let payload = data.range( + of: Self.payloadMarker, + in: data.startIndex.. + ) -> String? { + guard let marker = data.range( + of: Self.typeValueMarker, + in: range + ), let end = data[marker.upperBound.. String? { switch wireValue { case "CommandExecution": @@ -1031,6 +1141,11 @@ struct IncrementalRolloutTailSource { } } +private enum TokenActivityHeaderDisposition { + case decode + case skip +} + private struct RolloutWire: Decodable { let timestamp: String? let ordinal: UInt64? @@ -1041,6 +1156,7 @@ private struct RolloutWire: Decodable { let type: String? let id: String? let parentThreadId: String? + let cwd: String? let cliVersion: String? let historyMode: String? let agentRole: String? diff --git a/Sources/CodexLimits/UsageIntelligenceEngine.swift b/Sources/CodexLimits/UsageIntelligenceEngine.swift index 3535d5b..301a860 100644 --- a/Sources/CodexLimits/UsageIntelligenceEngine.swift +++ b/Sources/CodexLimits/UsageIntelligenceEngine.swift @@ -396,6 +396,8 @@ struct UsageIntelligenceInput: Equatable, Sendable { let localActivityFacts: [LocalActivityFact] let localActivityHistoryFacts: [LocalActivityFact] let localActivityObservation: LocalActivityObservation + let precomputedLocalTokenActivity: LocalTokenActivitySnapshot? + let precomputedLocalFilterOptions: UsageReceiptFilterOptions? let localTaskProjections: [ThreadProjection] let localActivityContentRevision: UInt64? let reusableLocalAggregates: LocalAggregateCache? @@ -417,6 +419,8 @@ struct UsageIntelligenceInput: Equatable, Sendable { localActivityObservation: LocalActivityObservation = .unavailable( "Codex local records are unavailable" ), + precomputedLocalTokenActivity: LocalTokenActivitySnapshot? = nil, + precomputedLocalFilterOptions: UsageReceiptFilterOptions? = nil, localTaskProjections: [ThreadProjection] = [], localActivityContentRevision: UInt64? = nil, reusableLocalAggregates: LocalAggregateCache? = nil, @@ -436,6 +440,8 @@ struct UsageIntelligenceInput: Equatable, Sendable { self.localActivityHistoryFacts = localActivityHistoryFacts ?? localActivityFacts self.localActivityObservation = localActivityObservation + self.precomputedLocalTokenActivity = precomputedLocalTokenActivity + self.precomputedLocalFilterOptions = precomputedLocalFilterOptions self.localTaskProjections = localTaskProjections self.localActivityContentRevision = localActivityContentRevision self.reusableLocalAggregates = reusableLocalAggregates @@ -559,6 +565,7 @@ struct UsageReaderSnapshot: Equatable, Sendable { let bankedResets: BankedResetSummary? let accountTokenActivity: AccountTokenActivitySnapshot let localTokenActivity: LocalTokenActivitySnapshot + let localFilterOptions: UsageReceiptFilterOptions? let usagePerToken: UsagePerTokenSnapshot let usageReceipts: UsageReceiptSnapshot let activityTimeline: ActivityTimelineSnapshot @@ -739,7 +746,12 @@ enum UsageIntelligenceEngine { samples: currentSamples ) let localTokenActivity: LocalTokenActivitySnapshot - if let interval = tokenActivityInterval( + if let precomputed = input.precomputedLocalTokenActivity { + localTokenActivity = precomputed.updating( + interval: precomputed.interval, + observation: input.localActivityObservation + ) + } else if let interval = tokenActivityInterval( account: input.account, accountActivity: accountTokenActivity, accountEpochStartedAt: input.accountEpochStartedAt @@ -980,6 +992,7 @@ enum UsageIntelligenceEngine { ), accountTokenActivity: accountTokenActivity, localTokenActivity: localTokenActivity, + localFilterOptions: input.precomputedLocalFilterOptions, usagePerToken: usagePerToken, usageReceipts: usageReceipts, activityTimeline: activityTimeline, diff --git a/Sources/CodexLimits/UsageMonitor.swift b/Sources/CodexLimits/UsageMonitor.swift index 466a32f..9499a62 100644 --- a/Sources/CodexLimits/UsageMonitor.swift +++ b/Sources/CodexLimits/UsageMonitor.swift @@ -12,6 +12,18 @@ enum SafetyBufferPolicy { } } +private struct StoredTokenQuery: Equatable { + let timeRange: AnalyticsTimeRange + let visibleRange: DateInterval? + let filters: WorkspaceFilters + + init(_ exploration: AnalyticsExplorationState) { + timeRange = exploration.timeRange + visibleRange = exploration.visibleRange + filters = exploration.filters + } +} + @MainActor final class UsageMonitor: ObservableObject { static let safetyBufferKey = "safetyBuffer" @@ -76,12 +88,17 @@ final class UsageMonitor: ObservableObject { private var localActivityCollection = LocalActivityCollection.unavailable( "Codex local records are unavailable" ) + private var precomputedLocalTokenActivity: LocalTokenActivitySnapshot? + private var precomputedLocalFilterOptions: UsageReceiptFilterOptions? + private var analyticsExploration = AnalyticsExplorationState.initial private var evaluationGeneration: UInt64 = 0 private var evaluationTask: Task? private var localImportGeneration: UInt64 = 0 + private var localLoadTask: Task? private var localImportTask: Task? private var localAnalyticsVisible = false private var localAnalyticsNeedsLoad = false + private var loadedStoredTokenQuery: StoredTokenQuery? convenience init() { self.init( @@ -120,6 +137,9 @@ final class UsageMonitor: ObservableObject { self.evaluateUsage = evaluateUsage self.localActivityCollector = localActivityCollector self.codexAssistedHistory = codexAssistedHistory + analyticsExploration = AnalyticsWorkspaceStore.restoredState( + from: defaults + ) let storedSafetyBuffer = defaults.object( forKey: Self.safetyBufferKey ) as? Double @@ -233,8 +253,11 @@ final class UsageMonitor: ObservableObject { guard !isRefreshing else { return } isRefreshing = true defer { isRefreshing = false } - if includeLocalActivity { - cancelLocalImport() + let refreshesVisibleStoredActivity = localAnalyticsVisible + if includeLocalActivity + || localAnalyticsNeedsLoad + || refreshesVisibleStoredActivity { + cancelLocalActivityWork() } await restoreHistoryIfAvailable() @@ -250,7 +273,9 @@ final class UsageMonitor: ObservableObject { accountSnapshot = result.snapshot sourceState = .available if localAnalyticsVisible, - includeLocalActivity || localAnalyticsNeedsLoad { + includeLocalActivity + || localAnalyticsNeedsLoad + || refreshesVisibleStoredActivity { await refreshLocalActivity( for: result.snapshot, observedAt: result.snapshot.fetchedAt, @@ -308,7 +333,9 @@ final class UsageMonitor: ObservableObject { accountSnapshot = newSnapshot sourceState = .available if localAnalyticsVisible, - includeLocalActivity || localAnalyticsNeedsLoad { + includeLocalActivity + || localAnalyticsNeedsLoad + || refreshesVisibleStoredActivity { await refreshLocalActivity( for: newSnapshot, observedAt: newSnapshot.fetchedAt @@ -341,8 +368,15 @@ final class UsageMonitor: ObservableObject { } } - func setLocalAnalyticsVisible(_ isVisible: Bool) async { + func setLocalAnalyticsVisible( + _ isVisible: Bool, + exploration: AnalyticsExplorationState? = nil + ) async { + if let exploration { + analyticsExploration = exploration + } if isVisible { + guard analyticsExploration.usesStoredTokenActivity else { return } if !localAnalyticsVisible { localAnalyticsVisible = true localAnalyticsNeedsLoad = true @@ -354,7 +388,7 @@ final class UsageMonitor: ObservableObject { } localAnalyticsVisible = false localAnalyticsNeedsLoad = false - cancelLocalImport() + cancelLocalActivityWork() localActivityCollection = .unavailable( "Codex local records are unavailable" ) @@ -371,6 +405,7 @@ final class UsageMonitor: ObservableObject { let accountSnapshot else { return } + let generation = localImportGeneration let identityVerified = sourceState == .available && historyAccountIdentity != nil await refreshLocalActivity( @@ -378,6 +413,10 @@ final class UsageMonitor: ObservableObject { observedAt: identityVerified ? accountSnapshot.fetchedAt : Date(), identityVerified: identityVerified ) + guard generation == localImportGeneration, + localAnalyticsVisible else { + return + } _ = await recalculate() } @@ -525,7 +564,8 @@ final class UsageMonitor: ObservableObject { defaults.set(planType, forKey: Self.historyPlanTypeKey) } guard partition != historyPartition else { return } - cancelLocalImport() + cancelLocalActivityWork() + localAnalyticsNeedsLoad = localAnalyticsVisible historyPartition = partition historyConnectionActive = false if let data = try? JSONEncoder().encode(partition) { @@ -593,7 +633,7 @@ final class UsageMonitor: ObservableObject { guard !isUpdatingHistory else { return } isUpdatingHistory = true defer { isUpdatingHistory = false } - cancelLocalImport() + cancelLocalActivityWork() let localDeletedAt = Date() NotificationCenter.default.post( name: .codexAssistedHistoryDeleted, @@ -655,7 +695,7 @@ final class UsageMonitor: ObservableObject { guard !isUpdatingHistory else { return } isUpdatingHistory = true defer { isUpdatingHistory = false } - cancelLocalImport() + cancelLocalActivityWork() let assistedDeletionCutoff = defaults.object( forKey: Self.localHistoryDeletionCutoffKey ) as? Date @@ -745,7 +785,7 @@ final class UsageMonitor: ObservableObject { guard !isUpdatingHistory, canRebuildAvailableHistory else { return } isUpdatingHistory = true defer { isUpdatingHistory = false } - cancelLocalImport() + cancelLocalActivityWork() do { let result = try await fetchUsage() guard let account = result.account else { @@ -844,6 +884,7 @@ final class UsageMonitor: ObservableObject { let buffer = SafetyBufferPolicy.normalized( safetyBuffer ?? storedBuffer ) + let usesStoredTokenActivity = precomputedLocalTokenActivity != nil return UsageIntelligenceInput( account: accountSnapshot, samples: historyMatchesCurrentSnapshot ? samples : [], @@ -856,14 +897,18 @@ final class UsageMonitor: ObservableObject { localActivityFacts: localActivityCollection.facts, localActivityHistoryFacts: localActivityCollection.facts, localActivityObservation: localActivityCollection.observation, + precomputedLocalTokenActivity: precomputedLocalTokenActivity, + precomputedLocalFilterOptions: precomputedLocalFilterOptions, localTaskProjections: localActivityCollection.projections, - localActivityContentRevision: - localActivityCollection.contentRevision, - reusableLocalAggregates: - readerSnapshot.reusableLocalAggregates, + localActivityContentRevision: usesStoredTokenActivity + ? nil + : localActivityCollection.contentRevision, + reusableLocalAggregates: usesStoredTokenActivity + ? nil + : readerSnapshot.reusableLocalAggregates, analyticsExploration: analyticsExploration - ?? AnalyticsWorkspaceStore.restoredState(from: defaults), + ?? self.analyticsExploration, insightDispositions: insightDispositions ?? AnalyticsWorkspaceStore @@ -902,6 +947,7 @@ final class UsageMonitor: ObservableObject { exploration: AnalyticsExplorationState, dispositions: [String: InsightDisposition] ) { + analyticsExploration = exploration let input = DeterministicInsightInput( reader: readerSnapshot, exploration: exploration @@ -910,6 +956,23 @@ final class UsageMonitor: ObservableObject { input, dispositions: dispositions ) + if localAnalyticsVisible, !exploration.usesLocalActivity { + cancelLocalActivityWork() + localAnalyticsNeedsLoad = false + localActivityCollection = .unavailable( + "Codex local records are unavailable" + ) + } + if localAnalyticsVisible, exploration.usesLocalActivity { + if loadedStoredTokenQuery != StoredTokenQuery(exploration) + || localAnalyticsNeedsLoad { + scheduleLocalActivityReload( + exploration: exploration, + dispositions: dispositions + ) + return + } + } guard evaluationTask != nil else { return } let pending = beginRecalculation( analyticsExploration: exploration, @@ -923,6 +986,79 @@ final class UsageMonitor: ObservableObject { } } + private func scheduleLocalActivityReload( + exploration: AnalyticsExplorationState, + dispositions: [String: InsightDisposition] + ) { + cancelLocalActivityWork() + evaluationGeneration &+= 1 + evaluationTask?.cancel() + localAnalyticsNeedsLoad = true + localActivityCollection = .unavailable( + "Codex local records are unavailable" + ) + let generation = localImportGeneration + let task = Task { [weak self] in + guard let self else { return } + while isRefreshing, !Task.isCancelled { + try? await Task.sleep(for: .milliseconds(50)) + } + guard !Task.isCancelled, + generation == localImportGeneration, + localAnalyticsVisible, + let accountSnapshot else { + return + } + let identityVerified = sourceState == .available + && historyAccountIdentity != nil + await refreshLocalActivity( + for: accountSnapshot, + observedAt: identityVerified + ? accountSnapshot.fetchedAt + : Date(), + identityVerified: identityVerified + ) + guard !Task.isCancelled, + generation == localImportGeneration, + localAnalyticsVisible, + analyticsExploration == exploration else { + return + } + let pending = beginRecalculation( + analyticsExploration: exploration, + insightDispositions: dispositions + ) + if await finishRecalculation(pending) { + await reconcileResetReminder() + } + if generation == localImportGeneration { + localLoadTask = nil + } + } + localLoadTask = task + } + + private func storedTokenActivityInterval( + weeklyInterval: DateInterval, + observedAt: Date, + exploration: AnalyticsExplorationState + ) -> DateInterval { + guard exploration.timeRange != .currentWindow else { + return weeklyInterval + } + let bounds = DateInterval( + start: weeklyInterval.end.addingTimeInterval(-84 * 86_400), + end: weeklyInterval.end + ) + if exploration.timeRange == .selected { + return bounds + } + return exploration.timeRange.interval( + within: bounds, + endingAt: min(observedAt, bounds.end) + ) + } + private func reconcileResetReminder() async { await resetReminderCoordinator.reconcile( target: resetReminderTarget() @@ -956,9 +1092,9 @@ final class UsageMonitor: ObservableObject { observedAt: Date, identityVerified: Bool = true ) async { - cancelLocalImport() let generation = localImportGeneration localAnalyticsNeedsLoad = false + let exploration = analyticsExploration guard let interval = UsageIntelligenceEngine.tokenActivityInterval( account: snapshot, samples: historyMatchesCurrentSnapshot ? samples : [], @@ -980,9 +1116,27 @@ final class UsageMonitor: ObservableObject { localActivityCollection = .unavailable( "Codex local records are unavailable" ) - let collection = await localActivityCollector.refresh( - interval: interval, - observedAt: observedAt + precomputedLocalTokenActivity = nil + precomputedLocalFilterOptions = nil + loadedStoredTokenQuery = nil + let storedQuery = StoredTokenQuery(exploration) + let storedInterval = storedTokenActivityInterval( + weeklyInterval: interval, + observedAt: observedAt, + exploration: exploration + ) + + let stored = await localActivityCollector + .refreshStoredTokenActivity( + interval: storedInterval, + filters: exploration.filters, + observedAt: observedAt + ) + guard generation == localImportGeneration else { return } + applyStoredTokenActivity( + stored, + identityVerified: identityVerified, + query: storedQuery ) guard generation == localImportGeneration else { return } let deletionPending = @@ -991,20 +1145,16 @@ final class UsageMonitor: ObservableObject { if deletionPending { historyDeletionStatus = .pendingLocal } - localActivityCollection = identityVerified - ? collection - : collection.loweringCoverage( - "Codex account identity could not be checked" - ) - let importPending = await localActivityCollector.hasPendingImport() - guard generation == localImportGeneration, importPending else { + guard generation == localImportGeneration, stored.importPending else { return } localImportTask = Task(priority: .background) { [weak self] in try? await Task.sleep(for: .milliseconds(500)) - await self?.continueLocalActivityImport( + await self?.continueStoredTokenActivityImport( with: localActivityCollector, - interval: interval, + interval: storedInterval, + filters: exploration.filters, + query: storedQuery, observedAt: observedAt, identityVerified: identityVerified, generation: generation @@ -1012,32 +1162,57 @@ final class UsageMonitor: ObservableObject { } } - private func continueLocalActivityImport( + private func applyStoredTokenActivity( + _ stored: LocalStoredTokenActivityCollection, + identityVerified: Bool, + query: StoredTokenQuery + ) { + let collection = LocalActivityCollection( + facts: [], + projections: [], + observation: stored.observation, + bytesRead: stored.bytesRead, + contentRevision: 0 + ) + localActivityCollection = identityVerified + ? collection + : collection.loweringCoverage( + "Codex account identity could not be checked" + ) + precomputedLocalTokenActivity = stored.snapshot + precomputedLocalFilterOptions = stored.filterOptions + loadedStoredTokenQuery = query + } + + private func continueStoredTokenActivityImport( with collector: LocalActivityCollector, interval: DateInterval, + filters: WorkspaceFilters, + query: StoredTokenQuery, observedAt: Date, identityVerified: Bool, generation: UInt64 ) async { - var latest: LocalActivityCollection? - while !Task.isCancelled, await collector.hasPendingImport() { + var latest: LocalStoredTokenActivityCollection? + var importPending = true + while !Task.isCancelled, importPending { while isRefreshing, !Task.isCancelled { try? await Task.sleep(for: .milliseconds(100)) } guard generation == localImportGeneration else { return } - latest = nil - let collection = await collector.refresh( + let collection = await collector.refreshStoredTokenActivity( interval: interval, - observedAt: observedAt, - refreshMetadata: false + filters: filters, + observedAt: observedAt ) guard !Task.isCancelled, generation == localImportGeneration else { return } latest = collection - if await collector.hasPendingImport() { - try? await Task.sleep(for: .milliseconds(250)) + importPending = collection.importPending + if importPending { + try? await Task.sleep(for: .milliseconds(100)) } } guard let latest, @@ -1045,11 +1220,11 @@ final class UsageMonitor: ObservableObject { generation == localImportGeneration else { return } - localActivityCollection = identityVerified - ? latest - : latest.loweringCoverage( - "Codex account identity could not be checked" - ) + applyStoredTokenActivity( + latest, + identityVerified: identityVerified, + query: query + ) _ = await recalculate(now: observedAt) persist() if generation == localImportGeneration { @@ -1057,10 +1232,15 @@ final class UsageMonitor: ObservableObject { } } - private func cancelLocalImport() { + private func cancelLocalActivityWork() { localImportGeneration &+= 1 + localLoadTask?.cancel() + localLoadTask = nil localImportTask?.cancel() localImportTask = nil + precomputedLocalTokenActivity = nil + precomputedLocalFilterOptions = nil + loadedStoredTokenQuery = nil } private func persist() { diff --git a/Tests/CodexLimitsTests/AnalyticsWorkspaceTests.swift b/Tests/CodexLimitsTests/AnalyticsWorkspaceTests.swift index c5129e4..74afb7e 100644 --- a/Tests/CodexLimitsTests/AnalyticsWorkspaceTests.swift +++ b/Tests/CodexLimitsTests/AnalyticsWorkspaceTests.swift @@ -26,8 +26,8 @@ final class AnalyticsWorkspaceTests: XCTestCase { defer { defaults.removePersistentDomain(forName: suiteName) } let first = AnalyticsWorkspaceStore(defaults: defaults) - first.selectSection(.facts) - first.selectGraph(.concurrency) + first.selectSection(.graphs) + first.selectGraph(.tokenActivity) first.selectTimeRange(.threeDays) first.updateFilters( WorkspaceFilters( @@ -50,8 +50,8 @@ final class AnalyticsWorkspaceTests: XCTestCase { let restored = AnalyticsWorkspaceStore(defaults: defaults) - XCTAssertEqual(restored.state.section, .facts) - XCTAssertEqual(restored.state.graph, .concurrency) + XCTAssertEqual(restored.state.section, .graphs) + XCTAssertEqual(restored.state.graph, .tokenActivity) XCTAssertEqual(restored.state.timeRange, .selected) XCTAssertEqual(restored.state.filters.projectID, "codex-limits") XCTAssertEqual(restored.state.filters.taskTreeID, "task-42") @@ -76,6 +76,20 @@ final class AnalyticsWorkspaceTests: XCTestCase { ) } + func testRestoredDeferredViewFallsBackToTheCoreWorkspace() { + let suiteName = "AnalyticsWorkspaceTests-\(UUID().uuidString)" + let defaults = UserDefaults(suiteName: suiteName)! + defer { defaults.removePersistentDomain(forName: suiteName) } + let store = AnalyticsWorkspaceStore(defaults: defaults) + store.selectSection(.insights) + store.selectGraph(.concurrency) + + XCTAssertEqual( + AnalyticsWorkspaceStore(defaults: defaults).state, + .initial + ) + } + func testPinnedUsageBaselinePersistsUntilCleared() { let suiteName = "AnalyticsWorkspaceTests-pin-\(UUID().uuidString)" let defaults = UserDefaults(suiteName: suiteName)! @@ -149,15 +163,46 @@ final class AnalyticsWorkspaceTests: XCTestCase { XCTAssertFalse(AnalyticsGraph.tokenActivity.usesAccountScope) } - func testOnlyUsageRemainingSkipsLocalActivity() { + func testOnlyTokenActivityLoadsLocalActivity() { var state = AnalyticsExplorationState.initial XCTAssertFalse(state.usesLocalActivity) state.graph = .tokenActivity XCTAssertTrue(state.usesLocalActivity) - state.graph = .usageRemaining state.section = .facts - XCTAssertTrue(state.usesLocalActivity) + XCTAssertFalse(state.usesLocalActivity) + state.section = .graphs + state.graph = .usagePerToken + XCTAssertFalse(state.usesLocalActivity) + state.graph = .concurrency + XCTAssertFalse(state.usesLocalActivity) + state.section = .insights + XCTAssertFalse(state.usesLocalActivity) + } + + func testEveryTokenGraphUsesTheBoundedTokenStore() { + var state = AnalyticsExplorationState.initial + + state.graph = .tokenActivity + XCTAssertTrue(state.usesStoredTokenActivity) + + state.filters.model = "gpt-5.6" + XCTAssertTrue(state.usesStoredTokenActivity) + + state.filters = .all + state.timeRange = .fourWeeks + XCTAssertTrue(state.usesStoredTokenActivity) + + state.timeRange = .currentWindow + state.visibleRange = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 2_000) + ) + XCTAssertTrue(state.usesStoredTokenActivity) + + state.visibleRange = nil + state.section = .facts + XCTAssertFalse(state.usesStoredTokenActivity) } func testPresetRangeEndsAtLatestObservedTimeAndIsClampedToWindow() { diff --git a/Tests/CodexLimitsTests/LocalActivityCollectorTests.swift b/Tests/CodexLimitsTests/LocalActivityCollectorTests.swift index be5d2db..d903f45 100644 --- a/Tests/CodexLimitsTests/LocalActivityCollectorTests.swift +++ b/Tests/CodexLimitsTests/LocalActivityCollectorTests.swift @@ -1,4 +1,5 @@ import XCTest +import SQLite3 @testable import CodexLimits final class LocalActivityCollectorTests: XCTestCase { @@ -2384,6 +2385,1118 @@ final class LocalActivityCollectorTests: XCTestCase { 7 ) } + + func testStoredTokenActivityStreamsLargeLegacyFactsAcrossRestartAndCutoff() + async throws + { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let stateDirectory = fixture.root.appendingPathComponent("state") + let interval = try fixture.interval() + let seed = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await seed.selectPartition("stable-account") + _ = await seed.refresh( + interval: interval, + observedAt: interval.start.addingTimeInterval(60) + ) + let partition = stateDirectory.appendingPathComponent("stable-account") + let factsFile = try XCTUnwrap( + FileManager.default.contentsOfDirectory( + at: partition, + includingPropertiesForKeys: nil + ).first { $0.lastPathComponent.hasSuffix(".facts.jsonl") } + ) + let template = try XCTUnwrap( + String(contentsOf: factsFile, encoding: .utf8) + .split(separator: "\n") + .compactMap { + try? JSONDecoder().decode( + LocalActivityFact.self, + from: Data($0.utf8) + ) + } + .first { $0.key == .token } + ) + let formatter = ISO8601DateFormatter() + let cutoff = try XCTUnwrap( + formatter.date(from: "2026-07-28T10:01:30Z") + ) + let firstRetained = try XCTUnwrap( + formatter.date(from: "2026-07-28T10:02:00Z") + ) + let generation = template.source.sourceGeneration + let tokenLine: (String, Date, Int64, String) -> String = { + eventID, date, delta, unusedPayload in + """ + {"key":"token","availability":"available","numericDelta":\(delta),"reason":null,"eventID":"\(eventID)","eventTimestamp":"\(formatter.string(from: date))","source":{"sourceGeneration":\(generation)},"value":{"unusedPayload":"\(unusedPayload)"}} + """ + } + var legacy = Data() + legacy.append( + Data( + tokenLine( + "before-cutoff", + cutoff.addingTimeInterval(-30), + 99, + "" + ).utf8 + ) + ) + legacy.append(0x0A) + legacy.append( + Data( + tokenLine( + "first-retained", + firstRetained, + 500, + String(repeating: "x", count: 800_000) + ).utf8 + ) + ) + legacy.append(0x0A) + for index in 1 ... 9_999 { + legacy.append( + Data( + tokenLine( + "stored-\(index)", + firstRetained.addingTimeInterval( + TimeInterval(index) + ), + 1, + "" + ).utf8 + ) + ) + legacy.append(0x0A) + } + try legacy.write(to: factsFile) + struct TestDeletionMarker: Encodable { + let version: Int + let cutoff: Date + let cleanupPending: Bool + } + let marker = stateDirectory.deletingLastPathComponent() + .appendingPathComponent( + ".\(stateDirectory.lastPathComponent)-deletion.json" + ) + try JSONEncoder().encode( + TestDeletionMarker( + version: 1, + cutoff: cutoff, + cleanupPending: false + ) + ).write(to: marker) + + let first = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await first.selectPartition("stable-account") + let partial = await first.refreshStoredTokenActivity( + interval: interval + ) + + XCTAssertTrue(partial.importPending) + XCTAssertNil(partial.snapshot.tokens) + await first.selectPartition("other-account") + + let restarted = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await restarted.selectPartition("stable-account") + var completed = await restarted.refreshStoredTokenActivity( + interval: interval + ) + for _ in 0 ..< 5 where completed.importPending { + completed = await restarted.refreshStoredTokenActivity( + interval: interval + ) + } + + XCTAssertFalse(completed.importPending) + XCTAssertEqual(completed.snapshot.tokens, 9_999) + XCTAssertEqual(completed.snapshot.points.last?.tokens, 9_999) + XCTAssertLessThanOrEqual(completed.snapshot.points.count, 1_000) + } + + func testStoredTokenActivityRebuildsAnOlderDerivedSchema() async throws { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let stateDirectory = fixture.root.appendingPathComponent("state") + let interval = try fixture.interval() + let seed = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await seed.selectPartition("stable-account") + _ = await seed.refresh( + interval: interval, + observedAt: interval.start.addingTimeInterval(60) + ) + let database = stateDirectory + .appendingPathComponent("stable-account") + .appendingPathComponent("local-activity-v1.sqlite3") + var connection: OpaquePointer? + XCTAssertEqual(sqlite3_open(database.path, &connection), SQLITE_OK) + XCTAssertEqual( + sqlite3_exec( + connection, + "PRAGMA user_version = 1", + nil, + nil, + nil + ), + SQLITE_OK + ) + sqlite3_close_v2(connection) + + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await collector.selectPartition("stable-account") + var result = await collector.refreshStoredTokenActivity( + interval: interval + ) + while result.importPending { + result = await collector.refreshStoredTokenActivity( + interval: interval + ) + } + + XCTAssertEqual(result.snapshot.tokens, 500) + } + + func testStoredTokenActivitySkipsSavedSourcesOutsideTheRange() async throws { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let stateDirectory = fixture.root.appendingPathComponent("state") + let interval = try fixture.interval() + let seed = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await seed.selectPartition("stable-account") + _ = await seed.refresh(interval: interval) + let partition = stateDirectory.appendingPathComponent("stable-account") + let entries = try FileManager.default.contentsOfDirectory( + at: partition, + includingPropertiesForKeys: nil + ) + let metadata = try XCTUnwrap( + entries.first { + $0.pathExtension == "json" + && !$0.lastPathComponent.hasSuffix(".facts.jsonl") + } + ) + let facts = partition.appendingPathComponent( + metadata.deletingPathExtension().lastPathComponent + + ".facts.jsonl" + ) + let oldMetadata = partition.appendingPathComponent("old.json") + let oldFacts = partition.appendingPathComponent("old.facts.jsonl") + var object = try XCTUnwrap( + JSONSerialization.jsonObject(with: Data(contentsOf: metadata)) + as? [String: Any] + ) + object["activityStart"] = 0.0 + object["activityEnd"] = 1.0 + try JSONSerialization.data(withJSONObject: object) + .write(to: oldMetadata) + try FileManager.default.copyItem(at: facts, to: oldFacts) + + let collector = LocalActivityCollector( + rootDirectory: fixture.root.appendingPathComponent("empty"), + stateDirectory: stateDirectory + ) + await collector.selectPartition("stable-account") + var result = await collector.refreshStoredTokenActivity( + interval: interval + ) + while result.importPending { + result = await collector.refreshStoredTokenActivity( + interval: interval + ) + } + + let database = partition.appendingPathComponent( + "local-activity-v1.sqlite3" + ) + var connection: OpaquePointer? + XCTAssertEqual(sqlite3_open(database.path, &connection), SQLITE_OK) + defer { sqlite3_close_v2(connection) } + var statement: OpaquePointer? + XCTAssertEqual( + sqlite3_prepare_v2( + connection, + "SELECT COUNT(*) FROM source WHERE is_active = 1", + -1, + &statement, + nil + ), + SQLITE_OK + ) + defer { sqlite3_finalize(statement) } + XCTAssertEqual(sqlite3_step(statement), SQLITE_ROW) + XCTAssertEqual(sqlite3_column_int(statement, 0), 1) + } + + func testStoredTokenActivityKeepsPersistedMalformedCoverage() async throws { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + #"{"timestamp":"2026-07-28T10:01:30Z","type":"event_msg""#, + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let stateDirectory = fixture.root.appendingPathComponent("state") + let interval = try fixture.interval() + let seed = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await seed.selectPartition("stable-account") + let seeded = await seed.refresh(interval: interval) + XCTAssertEqual( + seeded.observation.reason, + "Some local diagnostic records could not be read" + ) + let emptyRoot = fixture.root.appendingPathComponent( + "empty-rollouts", + isDirectory: true + ) + try FileManager.default.createDirectory( + at: emptyRoot, + withIntermediateDirectories: true + ) + let restarted = LocalActivityCollector( + rootDirectory: emptyRoot, + stateDirectory: stateDirectory + ) + await restarted.selectPartition("stable-account") + + var stored = await restarted.refreshStoredTokenActivity( + interval: interval + ) + for _ in 0 ..< 10 where stored.importPending { + stored = await restarted.refreshStoredTokenActivity( + interval: interval + ) + } + + XCTAssertEqual(stored.snapshot.tokens, 500) + XCTAssertEqual(stored.snapshot.coverage, .low) + XCTAssertEqual( + stored.snapshot.reason, + "Some local diagnostic records could not be read" + ) + } + + func testStoredTokenActivityKeepsPersistedContinuityGap() async throws { + let fixture = try CollectorFixture() + let file = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let stateDirectory = fixture.root.appendingPathComponent("state") + let interval = try fixture.interval() + let seed = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await seed.selectPartition("stable-account") + _ = await seed.refresh( + interval: interval, + observedAt: interval.start.addingTimeInterval(60) + ) + var replacement = + fixture.session(threadID: "task-1", ordinal: 0) + "\n" + for ordinal in 1 ... 5_000 { + replacement += + #"{"timestamp":"2026-07-28T10:01:00.000Z","ordinal":\#(ordinal),"type":"compacted","payload":{}}"# + + "\n" + } + replacement += fixture.tokens( + total: 900, + ordinal: 5_001, + minute: 3 + ) + "\n" + try Data(replacement.utf8).write(to: file, options: .atomic) + let changed = await seed.refresh( + interval: interval, + observedAt: interval.start.addingTimeInterval(180) + ) + XCTAssertEqual( + changed.observation.reason, + "Local task record continuity changed" + ) + let emptyRoot = fixture.root.appendingPathComponent( + "empty-rollouts", + isDirectory: true + ) + try FileManager.default.createDirectory( + at: emptyRoot, + withIntermediateDirectories: true + ) + let restarted = LocalActivityCollector( + rootDirectory: emptyRoot, + stateDirectory: stateDirectory + ) + await restarted.selectPartition("stable-account") + + var stored = await restarted.refreshStoredTokenActivity( + interval: interval + ) + for _ in 0 ..< 10 where stored.importPending { + stored = await restarted.refreshStoredTokenActivity( + interval: interval + ) + } + + XCTAssertEqual(stored.snapshot.coverage, .low) + XCTAssertEqual( + stored.snapshot.reason, + "Local task record continuity changed" + ) + } + + func testStoredTokenActivityRetriesEveryFactAfterASQLiteWriteFailure() + async throws + { + let fixture = try CollectorFixture() + let file = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let stateDirectory = fixture.root.appendingPathComponent("state") + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await collector.selectPartition("stable-account") + let interval = try fixture.interval() + let first = await collector.refreshStoredTokenActivity( + interval: interval + ) + XCTAssertEqual(first.snapshot.tokens, 500) + let partition = stateDirectory.appendingPathComponent( + "stable-account", + isDirectory: true + ) + let database = partition.appendingPathComponent( + "local-activity-v1.sqlite3" + ) + var lock: OpaquePointer? + XCTAssertEqual(sqlite3_open(database.path, &lock), SQLITE_OK) + XCTAssertEqual( + sqlite3_exec(lock, "BEGIN IMMEDIATE", nil, nil, nil), + SQLITE_OK + ) + defer { + sqlite3_exec(lock, "ROLLBACK", nil, nil, nil) + sqlite3_close_v2(lock) + } + try fixture.append( + fixture.tokens(total: 800, ordinal: 3, minute: 3), + to: file + ) + let failed = await collector.refreshStoredTokenActivity( + interval: interval + ) + XCTAssertEqual(failed.snapshot.coverage, .unavailable) + XCTAssertEqual(sqlite3_exec(lock, "ROLLBACK", nil, nil, nil), SQLITE_OK) + XCTAssertEqual(sqlite3_close_v2(lock), SQLITE_OK) + lock = nil + try fixture.append( + fixture.tokens(total: 1_000, ordinal: 4, minute: 4), + to: file + ) + + let recovered = await collector.refreshStoredTokenActivity( + interval: interval + ) + let unchanged = await collector.refreshStoredTokenActivity( + interval: interval + ) + + XCTAssertEqual(recovered.snapshot.tokens, 900) + XCTAssertEqual(unchanged.snapshot.tokens, 900) + } + + func testCancelledStoredTokenRefreshStopsInsteadOfReportingAReadFailure() + async throws + { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let stateDirectory = fixture.root.appendingPathComponent("state") + let seed = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await seed.selectPartition("stable-account") + _ = await seed.refresh(interval: try fixture.interval()) + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await collector.selectPartition("stable-account") + let interval = try fixture.interval() + let refresh = Task { + while !Task.isCancelled { + await Task.yield() + } + return await collector.refreshStoredTokenActivity( + interval: interval + ) + } + refresh.cancel() + + let result = await refresh.value + + XCTAssertEqual(result.snapshot.coverage, .unavailable) + XCTAssertEqual( + result.snapshot.reason, + "Local activity read was cancelled" + ) + } + + func testStoredTokenActivityTailsOnlyNewRolloutBytesAfterRestart() + async throws + { + let fixture = try CollectorFixture() + let file = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let stateDirectory = fixture.root.appendingPathComponent("state") + let interval = try fixture.interval() + do { + let seed = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await seed.selectPartition("stable-account") + _ = await seed.refresh(interval: interval) + let importing = await seed.refreshStoredTokenActivity( + interval: interval + ) + XCTAssertTrue(importing.importPending) + let migrated = await seed.refreshStoredTokenActivity( + interval: interval + ) + XCTAssertEqual(migrated.snapshot.tokens, 500) + } + try fixture.append( + fixture.tokens(total: 800, ordinal: 3, minute: 3), + to: file + ) + let projectionRequests = CollectorProjectionRequests() + let restarted = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory, + projectionSource: ReadOnlyThreadProjectionSource { request in + await projectionRequests.record(request) + throw CocoaError(.fileReadUnknown) + } + ) + await restarted.selectPartition("stable-account") + + let updated = await restarted.refreshStoredTokenActivity( + interval: interval + ) + let unchanged = await restarted.refreshStoredTokenActivity( + interval: interval + ) + let requests = await projectionRequests.values + + XCTAssertEqual(updated.snapshot.tokens, 700) + XCTAssertGreaterThan(updated.bytesRead, 0) + XCTAssertEqual(unchanged.snapshot.tokens, 700) + XCTAssertEqual(unchanged.bytesRead, 0) + XCTAssertEqual(requests, []) + } + + func testStoredTokenActivityBootstrapsFromRolloutWithoutFullRefresh() + async throws + { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let projectionRequests = CollectorProjectionRequests() + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: fixture.root.appendingPathComponent("state"), + projectionSource: ReadOnlyThreadProjectionSource { request in + await projectionRequests.record(request) + throw CocoaError(.fileReadUnknown) + } + ) + await collector.selectPartition("stable-account") + + let result = await collector.refreshStoredTokenActivity( + interval: try fixture.interval() + ) + let requests = await projectionRequests.values + + XCTAssertEqual(result.snapshot.tokens, 500) + XCTAssertGreaterThan(result.bytesRead, 0) + XCTAssertEqual(requests, []) + } + + func testStoredTokenActivityImportsRawRolloutDirectlyAcrossRestart() + async throws + { + let fixture = try CollectorFixture() + var lines = [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1) + ] + lines.append( + contentsOf: (2 ..< 10_000).map { + #"{"timestamp":"2026-07-28T10:01:30.000Z","ordinal":\#($0),"type":"compacted","payload":{}}"# + } + ) + lines.append( + fixture.tokens(total: 600, ordinal: 10_000, minute: 2) + ) + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: lines + ) + let stateDirectory = fixture.root.appendingPathComponent("state") + let partition = stateDirectory.appendingPathComponent( + "stable-account" + ) + let interval = try fixture.interval() + + do { + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await collector.selectPartition("stable-account") + let first = await collector.refreshStoredTokenActivity( + interval: interval + ) + + XCTAssertTrue(first.importPending) + XCTAssertTrue( + FileManager.default.fileExists( + atPath: partition.appendingPathComponent( + "local-activity-v1.sqlite3" + ).path + ) + ) + XCTAssertEqual(try storedFactSidecars(in: partition), []) + } + + let restarted = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await restarted.selectPartition("stable-account") + var completed = await restarted.refreshStoredTokenActivity( + interval: interval + ) + for _ in 0 ..< 3 where completed.importPending { + completed = await restarted.refreshStoredTokenActivity( + interval: interval + ) + } + + XCTAssertFalse(completed.importPending) + XCTAssertEqual(completed.snapshot.tokens, 500) + XCTAssertEqual(try storedFactSidecars(in: partition), []) + } + + func testStoredTokenActivityCountsOnlyAChildTasksOwnCopiedHistory() + async throws + { + let fixture = try CollectorFixture() + let child = "019fb303-6225-7592-96dd-15504bb96fbd" + let parent = "019fa3d9-328e-7ca0-8d40-e61dec66d483" + _ = try fixture.rollout( + day: "2026/07/28", + threadID: child, + lines: [ + #"{"timestamp":"2026-07-28T10:00:00.000Z","ordinal":0,"type":"session_meta","payload":{"id":"\#(child)","parent_thread_id":"\#(parent)","cli_version":"0.145.0"}}"#, + #"{"timestamp":"2026-07-28T10:00:00.100Z","ordinal":1,"type":"session_meta","payload":{"id":"\#(parent)","cli_version":"0.145.0"}}"#, + #"{"timestamp":"2026-07-28T10:00:00.200Z","ordinal":2,"type":"event_msg","payload":{"type":"task_started","turn_id":"019fa3da-0000-7000-8000-000000000000"}}"#, + #"{"timestamp":"2026-07-28T10:00:00.300Z","ordinal":3,"type":"event_msg","payload":{"type":"task_started","turn_id":"e30f1c44-bcc4-4b16-ab4a-576ed6b5aa46"}}"#, + fixture.tokens(total: 100, ordinal: 4, minute: 1), + fixture.tokens(total: 600, ordinal: 5, minute: 2), + #"{"timestamp":"2026-07-28T10:03:00.000Z","ordinal":6,"type":"event_msg","payload":{"type":"task_started","turn_id":"019fb304-0000-7000-8000-000000000000"}}"#, + #"{"timestamp":"2026-07-28T10:04:00.000Z","ordinal":7,"type":"turn_context","payload":{"turn_id":"019fb304-0000-7000-8000-000000000000","model":"gpt-5.6","effort":"high"}}"#, + #"{"timestamp":"2026-07-28T10:05:00.000Z","ordinal":8,"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"total_tokens":700},"last_token_usage":{"total_tokens":100}}}}"#, + #"{"timestamp":"2026-07-28T10:06:00.000Z","ordinal":9,"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"total_tokens":900},"last_token_usage":{"total_tokens":200}}}}"# + ] + ) + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: fixture.root.appendingPathComponent("state") + ) + await collector.selectPartition("stable-account") + + let result = await collector.refreshStoredTokenActivity( + interval: try fixture.interval() + ) + + XCTAssertEqual(result.snapshot.tokens, 300) + XCTAssertEqual(result.filterOptions.models, ["gpt-5.6"]) + XCTAssertEqual(result.filterOptions.reasoningLevels, ["high"]) + } + + func testStoredTokenActivityDeduplicatesAResumedTaskCopy() + async throws + { + let fixture = try CollectorFixture() + let task = "019fb303-6225-7592-96dd-15504bb96fbd" + let turn = "019fb304-0000-7000-8000-000000000000" + let lines = [ + #"{"timestamp":"2026-07-28T10:00:00.000Z","type":"session_meta","payload":{"id":"\#(task)","cli_version":"0.145.0"}}"#, + #"{"timestamp":"2026-07-28T10:00:30.000Z","type":"turn_context","payload":{"turn_id":"\#(turn)","model":"gpt-5.6","effort":"high"}}"#, + #"{"timestamp":"2026-07-28T10:01:00.000Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"total_tokens":100}}}}"#, + #"{"timestamp":"2026-07-28T10:02:00.000Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"total_tokens":600}}}}"# + ] + _ = try fixture.rollout( + day: "2026/07/27", + threadID: task, + lines: lines + ) + _ = try fixture.rollout( + day: "2026/07/28", + threadID: task, + lines: lines + ) + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: fixture.root.appendingPathComponent("state") + ) + await collector.selectPartition("stable-account") + + let result = await collector.refreshStoredTokenActivity( + interval: try fixture.interval() + ) + + XCTAssertEqual(result.snapshot.tokens, 500) + } + + func testStoredTokenActivityBuildsProjectFiltersFromRolloutMetadata() + async throws + { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + """ + {"timestamp":"2026-07-28T10:00:00.000Z","ordinal":0,"type":"session_meta","payload":{"id":"task-1","cli_version":"0.145.0","cwd":"/work/client-project"}} + """, + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: fixture.root.appendingPathComponent("state") + ) + await collector.selectPartition("stable-account") + + let result = await collector.refreshStoredTokenActivity( + interval: try fixture.interval() + ) + + XCTAssertEqual(result.filterOptions.projects, ["client-project"]) + XCTAssertEqual(result.filterOptions.taskTrees, ["task-1"]) + } + + func testStoredTokenActivityAdvancesPastOneAcceptedLargeRecord() + async throws + { + let fixture = try CollectorFixture() + let padding = String(repeating: "x", count: 5 * 1_024 * 1_024) + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + #"{"timestamp":"2026-07-28T10:00:30.000Z","ordinal":1,"type":"compacted","payload":{"padding":"\#(padding)"}}"#, + fixture.tokens(total: 100, ordinal: 2, minute: 1), + fixture.tokens(total: 600, ordinal: 3, minute: 2) + ] + ) + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: fixture.root.appendingPathComponent("state") + ) + await collector.selectPartition("stable-account") + + let result = await collector.refreshStoredTokenActivity( + interval: try fixture.interval() + ) + + XCTAssertFalse(result.importPending) + XCTAssertEqual(result.snapshot.tokens, 500) + } + + func testStoredTokenActivityFindsAContinuingRolloutFromPreviousDay() + async throws + { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/27", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: fixture.root.appendingPathComponent("state") + ) + await collector.selectPartition("stable-account") + + let result = await collector.refreshStoredTokenActivity( + interval: try fixture.interval() + ) + + XCTAssertEqual(result.snapshot.tokens, 500) + } + + func testStoredTokenActivityFindsARecentlyUpdatedOlderRollout() + async throws + { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/20", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: fixture.root.appendingPathComponent("state") + ) + await collector.selectPartition("stable-account") + + let result = await collector.refreshStoredTokenActivity( + interval: try fixture.interval() + ) + + XCTAssertEqual(result.snapshot.tokens, 500) + } + + func testStoredTokenActivityDiscoversTheWholeSelectedRange() async throws { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/06/15", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: fixture.root.appendingPathComponent("state") + ) + await collector.selectPartition("stable-account") + + let result = await collector.refreshStoredTokenActivity( + interval: try fixture.interval( + start: "2026-06-01T00:00:00Z", + end: "2026-07-29T00:00:00Z" + ) + ) + + XCTAssertEqual(result.snapshot.tokens, 500) + } + + func testStoredTokenActivityKeepsLastActiveFactsAfterContinuityChange() + async throws + { + let fixture = try CollectorFixture() + let file = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let stateDirectory = fixture.root.appendingPathComponent("state") + let interval = try fixture.interval() + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await collector.selectPartition("stable-account") + let first = await collector.refreshStoredTokenActivity( + interval: interval + ) + try Data( + ( + fixture.session(threadID: "task-1", ordinal: 0) + + "\n" + + fixture.tokens(total: 900, ordinal: 3, minute: 3) + + "\n" + ).utf8 + ).write(to: file, options: .atomic) + + let changed = await collector.refreshStoredTokenActivity( + interval: interval + ) + let unchanged = await collector.refreshStoredTokenActivity( + interval: interval + ) + let restarted = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await restarted.selectPartition("stable-account") + let afterRestart = await restarted.refreshStoredTokenActivity( + interval: interval + ) + + XCTAssertEqual(first.snapshot.tokens, 500) + XCTAssertEqual(changed.snapshot.tokens, 500) + XCTAssertEqual(changed.snapshot.coverage, .low) + XCTAssertEqual( + changed.snapshot.reason, + "Local task record continuity changed" + ) + XCTAssertFalse(unchanged.importPending) + XCTAssertEqual(unchanged.bytesRead, 0) + XCTAssertFalse(afterRestart.importPending) + XCTAssertEqual(afterRestart.bytesRead, 0) + XCTAssertEqual(afterRestart.snapshot.tokens, 500) + } + + func testStoredTokenActivityRebuildsSQLiteFromLegacyAfterLiveAppend() + async throws + { + let fixture = try CollectorFixture() + let file = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let stateDirectory = fixture.root.appendingPathComponent("state") + let interval = try fixture.interval() + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await collector.selectPartition("stable-account") + _ = await collector.refreshStoredTokenActivity(interval: interval) + try fixture.append( + fixture.tokens(total: 800, ordinal: 3, minute: 3), + to: file + ) + let updated = await collector.refreshStoredTokenActivity( + interval: interval + ) + await collector.selectPartition("other-account") + let partition = stateDirectory.appendingPathComponent("stable-account") + for name in [ + "local-activity-v1.sqlite3", + "local-activity-v1.sqlite3-wal", + "local-activity-v1.sqlite3-shm", + "local-activity-v1.sqlite3-journal" + ] { + let file = partition.appendingPathComponent(name) + if FileManager.default.fileExists(atPath: file.path) { + try FileManager.default.removeItem(at: file) + } + } + let restarted = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: stateDirectory + ) + await restarted.selectPartition("stable-account") + + var rebuilt = await restarted.refreshStoredTokenActivity( + interval: interval + ) + while rebuilt.importPending { + rebuilt = await restarted.refreshStoredTokenActivity( + interval: interval + ) + } + + XCTAssertEqual(updated.snapshot.tokens, 700) + XCTAssertEqual(rebuilt.snapshot.tokens, 700) + } + + func testStoredTokenActivityNoChangeRefreshesReadNoRolloutBytes() + async throws + { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session(threadID: "task-1", ordinal: 0), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: fixture.root.appendingPathComponent("state") + ) + await collector.selectPartition("stable-account") + let interval = try fixture.interval() + _ = await collector.refreshStoredTokenActivity(interval: interval) + + for _ in 0 ..< 10 { + let unchanged = await collector.refreshStoredTokenActivity( + interval: interval + ) + XCTAssertEqual(unchanged.snapshot.tokens, 500) + XCTAssertEqual(unchanged.bytesRead, 0) + } + } + + func testStoredTokenActivityWithoutSavedOrRolloutDataIsUnavailable() + async throws + { + let fixture = try CollectorFixture() + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: fixture.root.appendingPathComponent("state") + ) + await collector.selectPartition("stable-account") + + let result = await collector.refreshStoredTokenActivity( + interval: try fixture.interval() + ) + + XCTAssertEqual(result.snapshot.coverage, .unavailable) + XCTAssertEqual( + result.snapshot.reason, + "No saved local activity is available" + ) + } + + func testStoredTokenActivityLowersCoverageForUncheckedCLIVersion() + async throws + { + let fixture = try CollectorFixture() + _ = try fixture.rollout( + day: "2026/07/28", + threadID: "task-1", + lines: [ + fixture.session( + threadID: "task-1", + ordinal: 0, + cliVersion: "0.146.0" + ), + fixture.tokens(total: 100, ordinal: 1, minute: 1), + fixture.tokens(total: 600, ordinal: 2, minute: 2) + ] + ) + let collector = LocalActivityCollector( + rootDirectory: fixture.root, + stateDirectory: fixture.root.appendingPathComponent("state") + ) + await collector.selectPartition("stable-account") + + let result = await collector.refreshStoredTokenActivity( + interval: try fixture.interval() + ) + + XCTAssertEqual(result.snapshot.tokens, 500) + XCTAssertEqual(result.snapshot.coverage, .low) + XCTAssertEqual( + result.snapshot.reason, + "This Codex CLI version has not been checked" + ) + } +} + +private func storedFactSidecars(in directory: URL) throws -> [String] { + try FileManager.default.contentsOfDirectory( + at: directory, + includingPropertiesForKeys: nil + ) + .filter { $0.lastPathComponent.hasSuffix(".facts.jsonl") } + .map(\.lastPathComponent) + .sorted() } private actor CollectorProjectionRequests { diff --git a/Tests/CodexLimitsTests/LocalActivityNormalizerTests.swift b/Tests/CodexLimitsTests/LocalActivityNormalizerTests.swift index e7ca512..081cc81 100644 --- a/Tests/CodexLimitsTests/LocalActivityNormalizerTests.swift +++ b/Tests/CodexLimitsTests/LocalActivityNormalizerTests.swift @@ -477,6 +477,7 @@ final class LocalActivityNormalizerTests: XCTestCase { eventType: nil, threadID: threadID, parentThreadID: nil, + cwd: nil, cliVersion: "0.145.0", historyMode: "paginated", agentRole: agent?.role, diff --git a/Tests/CodexLimitsTests/LocalActivityPerformanceTests.swift b/Tests/CodexLimitsTests/LocalActivityPerformanceTests.swift index 2803429..b90432c 100644 --- a/Tests/CodexLimitsTests/LocalActivityPerformanceTests.swift +++ b/Tests/CodexLimitsTests/LocalActivityPerformanceTests.swift @@ -370,6 +370,256 @@ final class LocalActivityPerformanceTests: XCTestCase { XCTAssertEqual(incremental.records.count, 1) } + func testBoundedReaderDoesNotRetainReadChunks() throws { + let directory = temporaryDirectory() + try FileManager.default.createDirectory( + at: directory, + withIntermediateDirectories: true + ) + let file = directory.appendingPathComponent("large-lines.jsonl") + try autoreleasepool { + guard FileManager.default.createFile( + atPath: file.path, + contents: nil + ) else { + throw CocoaError(.fileWriteUnknown) + } + let handle = try FileHandle(forWritingTo: file) + defer { try? handle.close() } + let prefix = Data(#"{"payload":""#.utf8) + let payload = Data(repeating: 0x78, count: 64 * 1_024 - 20) + let suffix = Data(#""}"#.utf8) + Data([0x0A]) + for _ in 0 ..< 1_024 { + try handle.write(contentsOf: prefix) + try handle.write(contentsOf: payload) + try handle.write(contentsOf: suffix) + } + } + let handle = try FileHandle(forReadingFrom: file) + defer { try? handle.close() } + let residentBefore = currentResidentBytes() + + let result = try BoundedJSONLReader.read( + handle: handle, + from: 0, + maximumRecordBytes: 128 * 1_024 + ) { _, _ in true } + + let residentAfter = currentResidentBytes() + let residentGrowth = residentAfter > residentBefore + ? residentAfter - residentBefore + : 0 + XCTAssertEqual(result.processedLineCount, 1_024) + XCTAssertLessThan(residentGrowth, 40 * 1_024 * 1_024) + } + + func testRealDataSQLitePerformanceContract() async throws { + let fixture = try realDataFixture() + let emptyRolloutRoot = temporaryDirectory() + let database = fixture.partitionDirectory.appendingPathComponent( + "local-activity-v1.sqlite3" + ) + XCTAssertFalse(FileManager.default.fileExists(atPath: database.path)) + let collector = fixture.collector( + rolloutRoot: emptyRolloutRoot + ) + await collector.selectPartition(fixture.partition) + let start = ProcessInfo.processInfo.systemUptime + var refreshes = 0 + var importResidentPeak: UInt64 = 0 + var result: LocalStoredTokenActivityCollection + repeat { + result = await collector.refreshStoredTokenActivity( + interval: fixture.interval + ) + refreshes += 1 + if result.importPending { + importResidentPeak = max( + importResidentPeak, + currentResidentBytes() + ) + } + if refreshes >= 200, result.importPending { + XCTFail("Import did not finish within 200 bounded refreshes") + return + } + } while result.importPending + var usage = rusage() + getrusage(RUSAGE_SELF, &usage) + let milliseconds = + (ProcessInfo.processInfo.systemUptime - start) * 1_000 + let databaseBytes = databaseFootprint(database) + XCTAssertEqual(result.snapshot.tokens, fixture.expectedTokens) + XCTAssertLessThan(milliseconds, 30_000) + XCTAssertLessThan(usage.ru_maxrss, 250 * 1_024 * 1_024) + XCTAssertLessThan(databaseBytes, 160 * 1_024 * 1_024) + + await collector.selectPartition("closed") + let reopened = fixture.collector( + rolloutRoot: emptyRolloutRoot + ) + await reopened.selectPartition(fixture.partition) + let coldStart = ProcessInfo.processInfo.systemUptime + let cold = await reopened.refreshStoredTokenActivity( + interval: fixture.interval + ) + let coldMilliseconds = + (ProcessInfo.processInfo.systemUptime - coldStart) * 1_000 + XCTAssertFalse(cold.importPending) + XCTAssertEqual(cold.bytesRead, 0) + XCTAssertEqual(cold.snapshot.tokens, fixture.expectedTokens) + XCTAssertLessThan(coldMilliseconds, 1_000) + + var interactionMaximum: Double = 0 + for days in [1, 3, 28, 84] { + let interval = DateInterval( + start: max( + fixture.interval.start, + fixture.interval.end.addingTimeInterval( + -Double(days) * 86_400 + ) + ), + end: fixture.interval.end + ) + let interactionStart = ProcessInfo.processInfo.systemUptime + let ranged = await reopened.refreshStoredTokenActivity( + interval: interval + ) + interactionMaximum = max( + interactionMaximum, + (ProcessInfo.processInfo.systemUptime - interactionStart) + * 1_000 + ) + XCTAssertFalse(ranged.importPending) + XCTAssertEqual(ranged.bytesRead, 0) + } + + _ = await reopened.refreshStoredTokenActivity( + interval: fixture.interval + ) + var filters: [WorkspaceFilters] = [] + if let project = cold.filterOptions.projects.first { + var value = WorkspaceFilters.all + value.projectID = project + filters.append(value) + } + if let task = cold.filterOptions.taskTrees.first { + var value = WorkspaceFilters.all + value.taskTreeID = task + filters.append(value) + } + if let model = cold.filterOptions.models.first { + var value = WorkspaceFilters.all + value.model = model + filters.append(value) + } + if let reasoning = cold.filterOptions.reasoningLevels.first { + var value = WorkspaceFilters.all + value.reasoning = reasoning + filters.append(value) + } + for filter in filters { + let interactionStart = ProcessInfo.processInfo.systemUptime + let filtered = await reopened.refreshStoredTokenActivity( + interval: fixture.interval, + filters: filter + ) + interactionMaximum = max( + interactionMaximum, + (ProcessInfo.processInfo.systemUptime - interactionStart) + * 1_000 + ) + XCTAssertFalse(filtered.importPending) + XCTAssertEqual(filtered.bytesRead, 0) + } + XCTAssertLessThan(interactionMaximum, 500) + + let first = await reopened.refreshStoredTokenActivity( + interval: fixture.interval + ) + let residentBefore = currentResidentBytes() + let repeatStart = ProcessInfo.processInfo.systemUptime + var last = first + for _ in 0 ..< 100 { + last = await reopened.refreshStoredTokenActivity( + interval: fixture.interval + ) + XCTAssertEqual(last.bytesRead, 0) + XCTAssertEqual(last.snapshot.tokens, first.snapshot.tokens) + } + let residentAfter = currentResidentBytes() + let residentGrowth = residentAfter > residentBefore + ? residentAfter - residentBefore + : 0 + let repeatMilliseconds = + (ProcessInfo.processInfo.systemUptime - repeatStart) * 1_000 + + print( + [ + "REAL_SQLITE_CONTRACT", + "import_refreshes=\(refreshes)", + String(format: "import_ms=%.3f", milliseconds), + "import_resident_peak_bytes=\(importResidentPeak)", + "peak_rss_bytes=\(usage.ru_maxrss)", + "database_bytes=\(databaseBytes)", + String(format: "cold_ms=%.3f", coldMilliseconds), + String( + format: "interaction_max_ms=%.3f", + interactionMaximum + ), + String(format: "repeat_100_ms=%.3f", repeatMilliseconds), + "resident_growth_bytes=\(residentGrowth)", + "tokens=\(last.snapshot.tokens ?? 0)" + ].joined(separator: " ") + ) + XCTAssertLessThan(residentGrowth, 8 * 1_024 * 1_024) + XCTAssertLessThan(repeatMilliseconds, 5_000) + } + + private struct RealDataFixture { + let stateRoot: URL + let partition: String + let interval: DateInterval + let expectedTokens: Int64 + + var partitionDirectory: URL { + stateRoot.appendingPathComponent(partition, isDirectory: true) + } + + func collector(rolloutRoot: URL) -> LocalActivityCollector { + LocalActivityCollector( + rootDirectory: rolloutRoot, + stateDirectory: stateRoot + ) + } + } + + private func realDataFixture() throws -> RealDataFixture { + let environment = ProcessInfo.processInfo.environment + guard let root = environment["CODEX_LIMITS_PERF_STATE_ROOT"], + let partition = environment["CODEX_LIMITS_PERF_PARTITION"], + let startValue = environment["CODEX_LIMITS_PERF_START"], + let endValue = environment["CODEX_LIMITS_PERF_END"], + let expectedValue = environment[ + "CODEX_LIMITS_PERF_EXPECTED_TOKENS" + ], + let start = TimeInterval(startValue), + let end = TimeInterval(endValue), + let expectedTokens = Int64(expectedValue), + start < end else { + throw XCTSkip("Real local activity fixture was not configured.") + } + return RealDataFixture( + stateRoot: URL(fileURLWithPath: root, isDirectory: true), + partition: partition, + interval: DateInterval( + start: Date(timeIntervalSince1970: start), + end: Date(timeIntervalSince1970: end) + ), + expectedTokens: expectedTokens + ) + } + private func currentResidentBytes() -> UInt64 { var info = mach_task_basic_info() var count = mach_msg_type_number_t( @@ -392,4 +642,16 @@ final class LocalActivityPerformanceTests: XCTestCase { guard result == KERN_SUCCESS else { return 0 } return UInt64(info.resident_size) } + + private func databaseFootprint(_ database: URL) -> UInt64 { + [database.path, database.path + "-wal", database.path + "-shm"] + .reduce(0) { result, path in + let size = ( + try? FileManager.default.attributesOfItem(atPath: path)[ + .size + ] as? NSNumber + )?.uint64Value ?? 0 + return result + size + } + } } diff --git a/Tests/CodexLimitsTests/LocalActivityStoreTests.swift b/Tests/CodexLimitsTests/LocalActivityStoreTests.swift new file mode 100644 index 0000000..01a6b37 --- /dev/null +++ b/Tests/CodexLimitsTests/LocalActivityStoreTests.swift @@ -0,0 +1,904 @@ +import XCTest +import SQLite3 +@testable import CodexLimits + +final class LocalActivityStoreTests: XCTestCase { + func testSchemaStoresACompactTokenIndexWithoutFactPayloads() throws { + let database = temporaryDatabase() + let store = try LocalActivityStore(fileURL: database) + store.close() + var connection: OpaquePointer? + XCTAssertEqual(sqlite3_open_v2(database.path, &connection, SQLITE_OPEN_READONLY, nil), SQLITE_OK) + defer { sqlite3_close_v2(connection) } + var statement: OpaquePointer? + XCTAssertEqual( + sqlite3_prepare_v2( + connection, + "SELECT name FROM pragma_table_info('token_fact') ORDER BY cid", + -1, + &statement, + nil + ), + SQLITE_OK + ) + defer { sqlite3_finalize(statement) } + var columns: [String] = [] + while sqlite3_step(statement) == SQLITE_ROW { + columns.append(String(cString: sqlite3_column_text(statement, 0))) + } + + XCTAssertTrue( + Set([ + "source_id", + "event_id", + "occurred_at", + "numeric_delta", + "task_id", + "effective_model", + "reasoning" + ]).isSubset(of: Set(columns)) + ) + XCTAssertFalse(columns.contains("payload")) + XCTAssertFalse(columns.contains("fact_json")) + + sqlite3_finalize(statement) + statement = nil + XCTAssertEqual( + sqlite3_prepare_v2( + connection, + """ + SELECT EXISTS( + SELECT 1 + FROM sqlite_master + WHERE type = 'table' AND name = 'token_event' + ) + """, + -1, + &statement, + nil + ), + SQLITE_OK + ) + XCTAssertEqual(sqlite3_step(statement), SQLITE_ROW) + XCTAssertEqual(sqlite3_column_int(statement, 0), 0) + } + + func testAppendedTokenActivitySurvivesReopen() throws { + let database = temporaryDatabase() + let interval = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 2_000) + ) + var store: LocalActivityStore? = try LocalActivityStore(fileURL: database) + try store?.append( + [tokenFact(id: "one", timestamp: 1_100, delta: 125)], + to: source(key: "rollout") + ) + store?.close() + store = try LocalActivityStore(fileURL: database) + + let snapshot = try store?.tokenActivity( + in: interval, + observation: .continuous( + sourceVersion: "0.145.0", + observedAt: interval.end + ) + ) + + XCTAssertEqual(snapshot?.tokens, 125) + XCTAssertEqual(snapshot?.points.map(\.tokens), [125]) + XCTAssertEqual(snapshot?.coverage, .high) + store?.close() + } + + func testTokenQueryUsesHalfOpenBoundsAndDeduplicatesEventsGlobally() throws { + let database = temporaryDatabase() + let store = try LocalActivityStore(fileURL: database) + let interval = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 2_000) + ) + try store.append( + [ + tokenFact(id: "shared", timestamp: 1_100, delta: 900), + tokenFact(id: "same-time", timestamp: 1_100, delta: 250), + tokenFact(id: "at-end", timestamp: 2_000, delta: 500) + ], + to: source(key: "second") + ) + try store.append( + [ + tokenFact(id: "before", timestamp: 999, delta: 50), + tokenFact(id: "at-start", timestamp: 1_000, delta: 100), + tokenFact(id: "shared", timestamp: 1_100, delta: 200) + ], + to: source(key: "first") + ) + try store.append( + [tokenFact(id: "at-start", timestamp: 1_000, delta: 100)], + to: source(key: "first") + ) + + let snapshot = try store.tokenActivity( + in: interval, + observation: .continuous( + sourceVersion: "0.145.0", + observedAt: interval.end + ) + ) + + XCTAssertEqual(snapshot.tokens, 550) + XCTAssertEqual(snapshot.points.map(\.tokens), [100, 550]) + XCTAssertEqual( + snapshot.points.map(\.date), + [ + Date(timeIntervalSince1970: 1_000), + Date(timeIntervalSince1970: 1_100) + ] + ) + } + + func testLaterFactForAnotherKeyOfTheSameEventIsNotDropped() throws { + let database = temporaryDatabase() + let store = try LocalActivityStore(fileURL: database) + let storedSource = source(key: "rollout") + try store.append( + [contextFact(id: "shared", timestamp: 1_100)], + to: storedSource + ) + try store.append( + [tokenFact(id: "shared", timestamp: 1_100, delta: 100)], + to: storedSource + ) + let interval = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 2_000) + ) + + XCTAssertEqual(try tokens(in: store, interval: interval), 100) + } + + func testUnavailableTokenFactDoesNotContributeToActivity() throws { + let database = temporaryDatabase() + let store = try LocalActivityStore(fileURL: database) + try store.append( + [ + tokenFact( + id: "unavailable", + timestamp: 1_100, + delta: 900, + availability: .unavailable + ), + tokenFact(id: "available", timestamp: 1_200, delta: 100) + ], + to: source(key: "rollout") + ) + let interval = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 2_000) + ) + + XCTAssertEqual(try tokens(in: store, interval: interval), 100) + } + + func testReplacementStaysHiddenUntilActivationAndCanBeRolledBack() throws { + let database = temporaryDatabase() + let store = try LocalActivityStore(fileURL: database) + let storedSource = source(key: "rollout", generation: 7) + let interval = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 2_000) + ) + try store.append( + [tokenFact(id: "old", timestamp: 1_100, delta: 100, generation: 7)], + to: storedSource + ) + let replacement = try store.beginReplacement(for: storedSource) + try store.appendReplacement( + [tokenFact(id: "new", timestamp: 1_100, delta: 400, generation: 7)], + to: storedSource, + storeGeneration: replacement + ) + + XCTAssertEqual( + try tokens(in: store, interval: interval), + 100 + ) + + try store.activateReplacement( + sourceKey: storedSource.key, + storeGeneration: replacement + ) + XCTAssertEqual( + try tokens(in: store, interval: interval), + 400 + ) + + let abandoned = try store.beginReplacement(for: storedSource) + try store.appendReplacement( + [tokenFact(id: "abandoned", timestamp: 1_100, delta: 900, generation: 7)], + to: storedSource, + storeGeneration: abandoned + ) + try store.rollbackReplacement( + sourceKey: storedSource.key, + storeGeneration: abandoned + ) + XCTAssertEqual( + try tokens(in: store, interval: interval), + 400 + ) + } + + func testReplacementFallsBackToDuplicateFromAnotherActiveSource() throws { + let database = temporaryDatabase() + let store = try LocalActivityStore(fileURL: database) + let interval = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 2_000) + ) + let preferred = source(key: "a") + try store.append( + [tokenFact(id: "shared", timestamp: 1_100, delta: 100)], + to: preferred + ) + try store.append( + [tokenFact(id: "shared", timestamp: 1_100, delta: 200)], + to: source(key: "b") + ) + let replacement = try store.beginReplacement(for: preferred) + try store.appendReplacement( + [tokenFact(id: "new", timestamp: 1_200, delta: 50)], + to: preferred, + storeGeneration: replacement + ) + + try store.activateReplacement( + sourceKey: preferred.key, + storeGeneration: replacement + ) + + XCTAssertEqual(try tokens(in: store, interval: interval), 250) + } + + func testIncompleteLegacyReplacementCannotBecomeActive() throws { + let database = temporaryDatabase() + let store = try LocalActivityStore(fileURL: database) + let active = source(key: "rollout", generation: 7) + try store.append( + [tokenFact(id: "old", timestamp: 1_100, delta: 100, generation: 7)], + to: active + ) + var importing = active + importing.legacyDevice = 4 + importing.legacyInode = 8 + importing.legacySize = 1_024 + importing.legacyModificationDate = Date(timeIntervalSince1970: 900) + importing.legacyOffset = 128 + importing.pendingFactJSON = Data("pending".utf8) + let replacement = try store.beginReplacement(for: importing) + try store.appendReplacement( + [tokenFact(id: "new", timestamp: 1_100, delta: 400, generation: 7)], + to: importing, + storeGeneration: replacement + ) + + XCTAssertThrowsError( + try store.activateReplacement( + sourceKey: importing.key, + storeGeneration: replacement + ) + ) { error in + XCTAssertEqual( + error as? LocalActivityStore.StoreError, + .invalidSource + ) + } + + let interval = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 2_000) + ) + XCTAssertEqual(try tokens(in: store, interval: interval), 100) + + importing.legacyOffset = importing.legacySize + importing.pendingFactJSON = nil + try store.appendReplacement( + [], + to: importing, + storeGeneration: replacement + ) + try store.activateReplacement( + sourceKey: importing.key, + storeGeneration: replacement + ) + XCTAssertEqual(try tokens(in: store, interval: interval), 400) + } + + func testTokenQueryKeepsExactTotalWithAtMostOneThousandPoints() throws { + let database = temporaryDatabase() + let store = try LocalActivityStore(fileURL: database) + let interval = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 3_000) + ) + try store.append( + (0..<1_501).map { index in + tokenFact( + id: "event-\(index)", + timestamp: 1_001 + TimeInterval(index), + delta: 1 + ) + }, + to: source(key: "rollout") + ) + + let snapshot = try store.tokenActivity( + in: interval, + observation: .continuous( + sourceVersion: "0.145.0", + observedAt: interval.end + ) + ) + + XCTAssertEqual(snapshot.tokens, 1_501) + XCTAssertLessThanOrEqual(snapshot.points.count, 1_000) + XCTAssertEqual(snapshot.points.last?.tokens, 1_501) + } + + func testUnboundedCounterLowersCoverageWithoutLosingKnownTokens() throws { + let database = temporaryDatabase() + let store = try LocalActivityStore(fileURL: database) + let interval = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 2_000) + ) + try store.append( + [ + tokenFact( + id: "baseline", + timestamp: 1_100, + delta: nil, + reason: "segment-baseline" + ), + tokenFact(id: "known", timestamp: 1_200, delta: 500) + ], + to: source(key: "rollout") + ) + + let snapshot = try store.tokenActivity( + in: interval, + observation: .continuous( + sourceVersion: "0.145.0", + observedAt: interval.end + ) + ) + + XCTAssertEqual(snapshot.tokens, 500) + XCTAssertEqual(snapshot.coverage, .low) + XCTAssertEqual( + snapshot.reason, + "Local token activity starts from an unbounded counter" + ) + } + + func testSQLiteTokenSnapshotMatchesTheInMemoryAggregator() throws { + let database = temporaryDatabase() + let store = try LocalActivityStore(fileURL: database) + let interval = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 2_000) + ) + let preferred = [ + tokenFact( + id: "baseline", + timestamp: 1_100, + delta: nil, + reason: "segment-baseline" + ), + tokenFact(id: "shared", timestamp: 1_200, delta: 200), + tokenFact(id: "known", timestamp: 1_300, delta: 300), + tokenFact(id: "at-end", timestamp: 2_000, delta: 900) + ] + try store.append(preferred, to: source(key: "a")) + try store.append( + [tokenFact(id: "shared", timestamp: 1_200, delta: 800)], + to: source(key: "b") + ) + let observation = LocalActivityObservation.continuous( + sourceVersion: "0.145.0", + observedAt: interval.end + ) + + let memory = LocalTokenActivityAggregator.evaluate( + facts: preferred, + interval: interval, + observation: observation + ) + let sqlite = try store.tokenActivity( + in: interval, + observation: observation + ) + + XCTAssertEqual(sqlite.tokens, memory.tokens) + XCTAssertEqual(sqlite.coverage, memory.coverage) + XCTAssertEqual(sqlite.reason, memory.reason) + XCTAssertEqual(sqlite.points, memory.points) + } + + func testFilteredTokenActivityAndOptionsUseStoredTaskMetadata() throws { + let store = try LocalActivityStore(fileURL: temporaryDatabase()) + let interval = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 2_000) + ) + let root = projection( + taskID: "root", + parentTaskID: nil, + project: "Codex Limits" + ) + let child = projection( + taskID: "child", + parentTaskID: "root", + project: nil + ) + try store.append( + [ + tokenFact( + id: "selected", + timestamp: 1_100, + delta: 200, + context: LocalActivityContext( + taskID: "child", + turnID: "turn-a", + agent: LocalAgentIdentity( + nickname: "Scout", + role: "explorer" + ), + effectiveModel: "gpt-5.6", + reasoning: "high" + ) + ), + tokenFact( + id: "other", + timestamp: 1_200, + delta: 300, + context: LocalActivityContext( + taskID: nil, + turnID: nil, + agent: nil, + effectiveModel: "gpt-5.5", + reasoning: "medium" + ) + ) + ], + to: source(key: "rollout"), + projections: [child, root] + ) + let observation = LocalActivityObservation.continuous( + sourceVersion: "0.145.0", + observedAt: interval.end + ) + + let filtered = try store.tokenActivity( + in: interval, + filters: WorkspaceFilters( + projectID: "Codex Limits", + taskTreeID: "root", + model: "gpt-5.6", + reasoning: "high" + ), + observation: observation + ) + let options = try store.filterOptions(in: interval) + + XCTAssertEqual(filtered.tokens, 200) + XCTAssertEqual(filtered.points.map(\.tokens), [200]) + XCTAssertEqual(options.projects, ["Codex Limits"]) + XCTAssertEqual(options.taskTrees, ["root"]) + XCTAssertEqual(options.models, ["gpt-5.5", "gpt-5.6"]) + XCTAssertEqual(options.reasoningLevels, ["high", "medium"]) + } + + func testTokenTotalOverflowFailsInsteadOfWrapping() throws { + let database = temporaryDatabase() + let store = try LocalActivityStore(fileURL: database) + let interval = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 2_000) + ) + try store.append( + [ + tokenFact( + id: "maximum", + timestamp: 1_100, + delta: .max + ), + tokenFact(id: "overflow", timestamp: 1_200, delta: 1) + ], + to: source(key: "rollout") + ) + + XCTAssertThrowsError( + try store.tokenActivity( + in: interval, + observation: .continuous( + sourceVersion: "0.145.0", + observedAt: interval.end + ) + ) + ) { error in + XCTAssertEqual( + error as? LocalActivityStore.StoreError, + .tokenTotalOverflow + ) + } + } + + func testInvalidTokenQueryIsRejected() throws { + let store = try LocalActivityStore(fileURL: temporaryDatabase()) + let instant = Date(timeIntervalSince1970: 1_000) + + XCTAssertThrowsError( + try store.tokenActivity( + in: DateInterval(start: instant, end: instant), + observation: .continuous( + sourceVersion: "0.145.0", + observedAt: instant + ) + ) + ) { error in + XCTAssertEqual( + error as? LocalActivityStore.StoreError, + .invalidQuery + ) + } + } + + func testImportingReplacementResumesAfterReopenWithItsCursorState() throws { + let database = temporaryDatabase() + var source = source(key: "rollout", generation: 9) + source.legacyDevice = 4 + source.legacyInode = 8 + source.legacyOffset = 1_024 + source.legacySize = 1_024 + source.legacyModificationDate = Date(timeIntervalSince1970: 900) + source.legacyOffset = 128 + source.pendingFactJSON = Data("pending".utf8) + var store: LocalActivityStore? = try LocalActivityStore(fileURL: database) + let storeGeneration = try XCTUnwrap( + store?.beginReplacement(for: source) + ) + try store?.appendReplacement( + [tokenFact(id: "first", timestamp: 1_100, delta: 100, generation: 9)], + to: source, + storeGeneration: storeGeneration + ) + store?.close() + + store = try LocalActivityStore(fileURL: database) + let resumed = try XCTUnwrap( + store?.resumableReplacement(matching: source) + ) + + XCTAssertEqual(resumed.storeGeneration, storeGeneration) + XCTAssertEqual(resumed.source.legacyOffset, 128) + XCTAssertEqual(resumed.source.pendingFactJSON, Data("pending".utf8)) + + var secondChunk = resumed.source + secondChunk.legacyOffset = 1_024 + secondChunk.pendingFactJSON = nil + try store?.appendReplacement( + [tokenFact(id: "second", timestamp: 1_200, delta: 200, generation: 9)], + to: secondChunk, + storeGeneration: resumed.storeGeneration + ) + try store?.activateReplacement( + sourceKey: source.key, + storeGeneration: resumed.storeGeneration + ) + let interval = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 2_000) + ) + XCTAssertEqual(try store.map { try tokens(in: $0, interval: interval) }, 300) + } + + func testHistoryCutoffMismatchFailsClosed() throws { + let database = temporaryDatabase() + let originalCutoff = Date(timeIntervalSince1970: 1_000) + let store = try LocalActivityStore( + fileURL: database, + historyCutoff: originalCutoff + ) + try store.append( + [tokenFact(id: "stored", timestamp: 1_100, delta: 100)], + to: source(key: "rollout") + ) + store.close() + + XCTAssertThrowsError( + try LocalActivityStore( + fileURL: database, + historyCutoff: Date(timeIntervalSince1970: 1_200) + ) + ) { error in + XCTAssertEqual( + error as? LocalActivityStore.StoreError, + .historyCutoffMismatch + ) + } + + XCTAssertNoThrow( + try LocalActivityStore( + fileURL: database, + historyCutoff: originalCutoff + ).close() + ) + } + + func testUnknownSchemaFailsClosed() throws { + let database = temporaryDatabase() + try FileManager.default.createDirectory( + at: database.deletingLastPathComponent(), + withIntermediateDirectories: true + ) + var connection: OpaquePointer? + XCTAssertEqual(sqlite3_open(database.path, &connection), SQLITE_OK) + XCTAssertEqual( + sqlite3_exec( + connection, + "PRAGMA user_version = 1", + nil, + nil, + nil + ), + SQLITE_OK + ) + sqlite3_close_v2(connection) + + XCTAssertThrowsError(try LocalActivityStore(fileURL: database)) { + XCTAssertEqual( + $0 as? LocalActivityStore.StoreError, + .schemaMismatch + ) + } + } + + func testActiveSourceMatchUsesTheLegacyFileSignature() throws { + let database = temporaryDatabase() + let store = try LocalActivityStore(fileURL: database) + var source = source(key: "rollout", generation: 3) + source.legacyDevice = 4 + source.legacyInode = 8 + source.legacySize = 1_024 + source.legacyModificationDate = Date(timeIntervalSince1970: 900) + source.legacyOffset = 128 + try store.append([], to: source) + + XCTAssertFalse(try store.hasActiveSource(matching: source)) + + source.legacyOffset = 1_024 + try store.append([], to: source) + XCTAssertTrue(try store.hasActiveSource(matching: source)) + XCTAssertTrue(try store.hasActiveSource(key: source.key)) + XCTAssertFalse(try store.hasActiveSource(key: "other")) + + source.legacyInode = 9 + XCTAssertFalse(try store.hasActiveSource(matching: source)) + } + + func testCancellationDuringAppendRollsBackProcessedFacts() async throws { + let database = temporaryDatabase() + let store = try LocalActivityStore(fileURL: database) + let storedSource = source(key: "rollout") + let facts = (0..<1_024).map { index in + tokenFact( + id: "event-\(index)", + timestamp: 1_100 + TimeInterval(index), + delta: 1 + ) + } + let enteredLoop = expectation(description: "append entered fact loop") + let continueAppend = DispatchSemaphore(value: 0) + store.testFactProgress = { index in + guard index == 256 else { return } + enteredLoop.fulfill() + continueAppend.wait() + } + let append = Task { + try store.append(facts, to: storedSource) + } + await fulfillment(of: [enteredLoop], timeout: 5) + append.cancel() + continueAppend.signal() + + do { + try await append.value + XCTFail("Cancelled append should fail") + } catch is CancellationError { + // Expected. + } + + let interval = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 3_000) + ) + XCTAssertEqual(try tokens(in: store, interval: interval), 0) + store.testFactProgress = nil + } + + func testCancelledTokenQueryFailsWithCancellationError() async throws { + let database = temporaryDatabase() + let store = try LocalActivityStore(fileURL: database) + try store.append( + [tokenFact(id: "stored", timestamp: 1_100, delta: 100)], + to: source(key: "rollout") + ) + let interval = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 2_000) + ) + let query = Task { + while !Task.isCancelled { + await Task.yield() + } + return try store.tokenActivity( + in: interval, + observation: .continuous( + sourceVersion: "0.145.0", + observedAt: interval.end + ) + ) + } + query.cancel() + + do { + _ = try await query.value + XCTFail("Cancelled query should fail") + } catch is CancellationError { + // Expected. + } + } + + func testCancelledFilterQueryFailsWithCancellationError() async throws { + let store = try LocalActivityStore(fileURL: temporaryDatabase()) + let interval = DateInterval( + start: Date(timeIntervalSince1970: 1_000), + end: Date(timeIntervalSince1970: 2_000) + ) + let query = Task { + while !Task.isCancelled { + await Task.yield() + } + return try store.filterOptions(in: interval) + } + query.cancel() + + do { + _ = try await query.value + XCTFail("Cancelled query should fail") + } catch is CancellationError { + // Expected. + } + } + + private func temporaryDatabase() -> URL { + FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString, isDirectory: true) + .appendingPathComponent("local-activity-v1.sqlite3") + } + + private func source( + key: String, + generation: UInt64 = 1 + ) -> LocalActivityStore.Source { + LocalActivityStore.Source( + key: key, + sourceGeneration: generation + ) + } + + private func tokenFact( + id: String, + timestamp: TimeInterval, + delta: Int64?, + reason: String? = nil, + generation: UInt64 = 1, + availability: LocalActivityAvailability = .available, + context: LocalActivityContext? = nil + ) -> LocalActivityFact { + LocalActivityFact( + key: .token, + availability: availability, + value: .tokens( + LocalTokenUsage( + inputTokens: 0, + cachedInputTokens: 0, + cacheWriteInputTokens: 0, + outputTokens: 0, + reasoningOutputTokens: 0, + totalTokens: delta ?? 0 + ) + ), + numericDelta: delta, + tokenSegment: 0, + reason: reason, + eventID: id, + eventTimestamp: ISO8601DateFormatter().string( + from: Date(timeIntervalSince1970: timestamp) + ), + source: LocalActivitySourceMetadata( + source: .rolloutJSONL, + sourceVersion: "0.145.0", + schemaVersion: "rollout-v1", + sourceGeneration: generation, + historyMode: nil, + observedAt: Date(timeIntervalSince1970: timestamp) + ), + context: context + ) + } + + private func projection( + taskID: String, + parentTaskID: String?, + project: String? + ) -> ThreadProjection { + ThreadProjection( + taskID: taskID, + parentTaskID: parentTaskID, + projectLabel: project, + rolloutFileURL: nil, + createdAt: nil, + updatedAt: Date(timeIntervalSince1970: 1_500), + source: LocalActivitySourceMetadata( + source: .appServerThreadList, + sourceVersion: "0.145.0", + schemaVersion: "app-server-v1", + sourceGeneration: 0, + historyMode: nil, + observedAt: Date(timeIntervalSince1970: 1_500) + ) + ) + } + + private func contextFact( + id: String, + timestamp: TimeInterval + ) -> LocalActivityFact { + LocalActivityFact( + key: .context, + availability: .available, + value: nil, + numericDelta: nil, + tokenSegment: nil, + reason: nil, + eventID: id, + eventTimestamp: ISO8601DateFormatter().string( + from: Date(timeIntervalSince1970: timestamp) + ), + source: LocalActivitySourceMetadata( + source: .rolloutJSONL, + sourceVersion: "0.145.0", + schemaVersion: "rollout-v1", + sourceGeneration: 1, + historyMode: nil, + observedAt: Date(timeIntervalSince1970: timestamp) + ) + ) + } + + private func tokens( + in store: LocalActivityStore, + interval: DateInterval + ) throws -> Int64? { + try store.tokenActivity( + in: interval, + observation: .continuous( + sourceVersion: "0.145.0", + observedAt: interval.end + ) + ).tokens + } +} diff --git a/Tests/CodexLimitsTests/LocalTokenActivityTests.swift b/Tests/CodexLimitsTests/LocalTokenActivityTests.swift index f9b9009..7d6048a 100644 --- a/Tests/CodexLimitsTests/LocalTokenActivityTests.swift +++ b/Tests/CodexLimitsTests/LocalTokenActivityTests.swift @@ -298,6 +298,61 @@ final class LocalTokenActivityTests: XCTestCase { XCTAssertNil(reader.localTokenActivity.accountComparison.numericPercent) } + func testReaderUsesAStoredTokenSnapshotInsteadOfScanningFacts() { + let now = Date(timeIntervalSince1970: 2_000_000) + let account = accountSnapshot( + fetchedAt: now, + lifetimeTokens: 1_600 + ) + let boundary = UsageSample( + observedAt: account.mainLimit!.window.startsAt, + remainingPercent: 100, + resetsAt: account.mainLimit!.window.resetsAt, + lifetimeTokens: 1_000 + ) + let interval = DateInterval( + start: boundary.observedAt, + end: now + ) + let stored = LocalTokenActivitySnapshot( + tokens: 900, + interval: interval, + coverage: .high, + reason: "Only local activity on this Mac is observed", + sourceVersion: "0.145.0", + observedAt: now, + points: [ + LocalTokenActivityPoint(date: now, tokens: 900) + ], + accountComparison: .unavailable + ) + + let reader = UsageIntelligenceEngine.evaluate( + UsageIntelligenceInput( + account: account, + samples: [boundary], + safetyBuffer: 3, + sourceState: .available, + now: now, + previousStatus: nil, + localActivityFacts: [ + tokenFact( + id: "legacy", + timestamp: now.timeIntervalSince1970 - 60, + delta: 100 + ) + ], + localActivityObservation: .continuous( + sourceVersion: "0.145.0", + observedAt: now + ), + precomputedLocalTokenActivity: stored + ) + ) + + XCTAssertEqual(reader.localTokenActivity, stored) + } + func testOffDeviceActivityDoesNotTurnTheAccountLocalGapIntoCoverage() { let now = Date(timeIntervalSince1970: 2_000_000) let account = accountSnapshot( diff --git a/Tests/CodexLimitsTests/RolloutTailSourceTests.swift b/Tests/CodexLimitsTests/RolloutTailSourceTests.swift index c302b65..5408e42 100644 --- a/Tests/CodexLimitsTests/RolloutTailSourceTests.swift +++ b/Tests/CodexLimitsTests/RolloutTailSourceTests.swift @@ -791,6 +791,32 @@ final class RolloutTailSourceTests: XCTestCase { XCTAssertGreaterThan(batch.cursor.byteOffset, 0) } + func testTokenActivityScopeSkipsLargeUnusedPayloadsBeforeDecode() throws { + let fixture = try TemporaryRollout() + let unused = String(repeating: "x", count: 1_000_000) + try fixture.append( + #"{"timestamp":"2026-07-27T10:00:00.000Z","ordinal":0,"type":"response_item","payload":{"type":"message","content":"\#(unused)","broken":not-json}}"# + + "\n" + + #"{"timestamp":"2026-07-27T10:00:10.000Z","ordinal":1,"type":"event_msg","payload":{"type":"item_completed","item":{"type":"CommandExecution","output":"\#(unused)","broken":not-json}}}"# + + "\n" + + #"{"timestamp":"2026-07-27T10:00:20.000Z","ordinal":2,"type":"compacted","payload":{"replacement_history":"\#(unused)","broken":not-json}}"# + + "\n" + + #"{"timestamp":"2026-07-27T10:01:00.000Z","ordinal":3,"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"total_tokens":200}}}}"# + + "\n" + ) + + let batch = try IncrementalRolloutTailSource().read( + fileURL: fixture.url, + cursor: nil, + observedAt: Date(timeIntervalSince1970: 100), + recordScope: .tokenActivity + ) + + XCTAssertEqual(batch.records.map(\.totalTokens), [200]) + XCTAssertEqual(batch.unsupportedRecordCount, 3) + XCTAssertEqual(batch.malformedRecordCount, 0) + } + func testCompactedRecordSkipsItsUnusedLargePayload() throws { let fixture = try TemporaryRollout() let sourceContent = String(repeating: "x", count: 1_000_000) diff --git a/Tests/CodexLimitsTests/UsageMonitorHistoryTests.swift b/Tests/CodexLimitsTests/UsageMonitorHistoryTests.swift index fb2156c..1af122c 100644 --- a/Tests/CodexLimitsTests/UsageMonitorHistoryTests.swift +++ b/Tests/CodexLimitsTests/UsageMonitorHistoryTests.swift @@ -663,11 +663,14 @@ final class UsageMonitorHistoryTests: XCTestCase { var exploration = AnalyticsExplorationState.initial exploration.section = .insights exploration.timeRange = .threeDays + let callsBeforePreferenceChange = evaluator.callCount monitor.analyticsPreferencesDidChange( exploration: exploration, dispositions: [:] ) + try await Task.sleep(for: .milliseconds(50)) + XCTAssertEqual(evaluator.callCount, callsBeforePreferenceChange) evaluator.release() await refresh.value let deadline = ContinuousClock.now + .seconds(2) @@ -678,6 +681,7 @@ final class UsageMonitorHistoryTests: XCTestCase { await Task.yield() } + XCTAssertEqual(evaluator.maximumConcurrentCalls, 1) XCTAssertEqual(evaluator.lastExploration, exploration) XCTAssertEqual(monitor.readerSnapshot.menuBarText, "80%") } @@ -803,8 +807,13 @@ final class UsageMonitorHistoryTests: XCTestCase { await firstMonitor.refresh() await firstMonitor.refresh() XCTAssertNil(firstMonitor.readerSnapshot.localTokenActivity.tokens) + var exploration = AnalyticsExplorationState.initial + exploration.graph = .tokenActivity - await firstMonitor.setLocalAnalyticsVisible(true) + await firstMonitor.setLocalAnalyticsVisible( + true, + exploration: exploration + ) XCTAssertEqual( firstMonitor.readerSnapshot.localTokenActivity.tokens, @@ -824,7 +833,10 @@ final class UsageMonitorHistoryTests: XCTestCase { await restarted.refresh() XCTAssertNil(restarted.readerSnapshot.localTokenActivity.tokens) - await restarted.setLocalAnalyticsVisible(true) + await restarted.setLocalAnalyticsVisible( + true, + exploration: exploration + ) XCTAssertEqual(restarted.readerSnapshot.localTokenActivity.tokens, 400) XCTAssertEqual(restarted.readerSnapshot.localTokenActivity.coverage, .low) @@ -834,7 +846,7 @@ final class UsageMonitorHistoryTests: XCTestCase { ) } - func testLocalCollectorReadsOnlyForVisibleAnalyticsAndManualRefresh() + func testDefaultViewDoesNotCreateLocalStoreAndTokenActivitySkipsFullScan() async throws { let suiteName = "UsageMonitorHistoryTests-\(UUID().uuidString)" @@ -850,6 +862,7 @@ final class UsageMonitorHistoryTests: XCTestCase { withIntermediateDirectories: true ) let requests = RequestCounter() + let stateRoot = root.appendingPathComponent("local-state") let fetches = DelayedFetchSource( makeFetchResult( identity: "user@example.com", @@ -863,7 +876,7 @@ final class UsageMonitorHistoryTests: XCTestCase { startsAutomatically: false, localActivityCollector: LocalActivityCollector( rootDirectory: localRoot, - stateDirectory: root.appendingPathComponent("local-state"), + stateDirectory: stateRoot, projectionSource: ReadOnlyThreadProjectionSource { _ in await requests.record() return Data( @@ -877,33 +890,47 @@ final class UsageMonitorHistoryTests: XCTestCase { await monitor.start() var requestCount = await requests.count XCTAssertEqual(requestCount, 0) + XCTAssertFalse(FileManager.default.fileExists(atPath: stateRoot.path)) let automatic = Task { @MainActor in await monitor.automaticRefresh() } try await Task.sleep(for: .milliseconds(10)) - await monitor.setLocalAnalyticsVisible(true) await automatic.value requestCount = await requests.count - XCTAssertEqual(requestCount, 1) - await monitor.setLocalAnalyticsVisible(true) + XCTAssertEqual(requestCount, 0) + XCTAssertFalse(FileManager.default.fileExists(atPath: stateRoot.path)) + + await monitor.refresh() requestCount = await requests.count - XCTAssertEqual(requestCount, 1) + XCTAssertEqual(requestCount, 0) + XCTAssertFalse(FileManager.default.fileExists(atPath: stateRoot.path)) + + var exploration = AnalyticsExplorationState.initial + exploration.graph = .tokenActivity + await monitor.setLocalAnalyticsVisible( + true, + exploration: exploration + ) + requestCount = await requests.count + XCTAssertEqual(requestCount, 0) await monitor.automaticRefresh() requestCount = await requests.count - XCTAssertEqual(requestCount, 1) + XCTAssertEqual(requestCount, 0) await monitor.refresh() requestCount = await requests.count - XCTAssertEqual(requestCount, 2) + XCTAssertEqual(requestCount, 0) await monitor.setLocalAnalyticsVisible(false) await monitor.refresh() requestCount = await requests.count - XCTAssertEqual(requestCount, 2) + XCTAssertEqual(requestCount, 0) } - func testHidingAnalyticsDiscardsAnInFlightLocalRead() async throws { + func testEveryTokenGraphUsesStoredActivityWithoutLoadingFullFacts() + async throws + { let suiteName = "UsageMonitorHistoryTests-\(UUID().uuidString)" let defaults = try XCTUnwrap(UserDefaults(suiteName: suiteName)) defer { defaults.removePersistentDomain(forName: suiteName) } @@ -912,44 +939,134 @@ final class UsageMonitorHistoryTests: XCTestCase { "rollouts", isDirectory: true ) + let rolloutDirectory = localRoot.appendingPathComponent( + "1970/01/22", + isDirectory: true + ) try FileManager.default.createDirectory( - at: localRoot, + at: rolloutDirectory, withIntermediateDirectories: true ) - let gate = ProjectionGate() - let fetchResult = makeFetchResult( - identity: "user@example.com", - fetchedAt: Date(timeIntervalSince1970: 1_700_000), - remaining: 90 + let fetchedAt = Date(timeIntervalSince1970: 1_850_000) + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [ + .withInternetDateTime, + .withFractionalSeconds + ] + let timestamps = [ + formatter.string(from: fetchedAt.addingTimeInterval(-120)), + formatter.string(from: fetchedAt.addingTimeInterval(-60)), + formatter.string(from: fetchedAt.addingTimeInterval(-1)) + ] + let lines = [ + #"{"timestamp":"\#(timestamps[0])","ordinal":0,"type":"session_meta","payload":{"id":"task-1","cli_version":"0.145.0"}}"#, + #"{"timestamp":"\#(timestamps[1])","ordinal":1,"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"total_tokens":100}}}}"#, + #"{"timestamp":"\#(timestamps[2])","ordinal":2,"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"total_tokens":500}}}}"# + ] + let rolloutFile = rolloutDirectory.appendingPathComponent( + "rollout-test.jsonl" + ) + try Data((lines.joined(separator: "\n") + "\n").utf8).write( + to: rolloutFile + ) + let requests = RequestCounter() + let collector = LocalActivityCollector( + rootDirectory: localRoot, + stateDirectory: root.appendingPathComponent("local-state"), + projectionSource: ReadOnlyThreadProjectionSource { _ in + await requests.record() + return Data( + #"{"result":{"data":[],"nextCursor":null}}"#.utf8 + ) + } ) + let source = FetchSequence([ + makeFetchResult( + identity: "user@example.com", + fetchedAt: Date(timeIntervalSince1970: 1_700_000), + remaining: 90 + ), + makeFetchResult( + identity: "user@example.com", + fetchedAt: fetchedAt, + remaining: 80 + ), + makeFetchResult( + identity: "user@example.com", + fetchedAt: fetchedAt.addingTimeInterval(10), + remaining: 79 + ) + ]) let monitor = UsageMonitor( defaults: defaults, historyDirectory: root.appendingPathComponent("history"), startsAutomatically: false, - localActivityCollector: LocalActivityCollector( - rootDirectory: localRoot, - stateDirectory: root.appendingPathComponent("local-state"), - projectionSource: ReadOnlyThreadProjectionSource { _ in - await gate.response() - } - ), - fetchUsage: { fetchResult } + localActivityCollector: collector, + fetchUsage: { try await source.next() } ) await monitor.refresh() + await monitor.refresh() + var exploration = AnalyticsExplorationState.initial + exploration.graph = .tokenActivity - let load = Task { @MainActor in - await monitor.setLocalAnalyticsVisible(true) + await monitor.setLocalAnalyticsVisible( + true, + exploration: exploration + ) + + for _ in 0 ..< 200 + where monitor.readerSnapshot.localTokenActivity.tokens == nil { + try await Task.sleep(for: .milliseconds(10)) } - await gate.waitUntilStarted() - let hide = Task { @MainActor in - await monitor.setLocalAnalyticsVisible(false) + XCTAssertEqual(monitor.readerSnapshot.localTokenActivity.tokens, 400) + var requestCount = await requests.count + XCTAssertEqual(requestCount, 0) + + let handle = try FileHandle(forWritingTo: rolloutFile) + try handle.seekToEnd() + try handle.write( + contentsOf: Data( + ( + #"{"timestamp":"\#(timestamps[2])","ordinal":3,"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"total_tokens":700}}}}"# + + "\n" + ).utf8 + ) + ) + try handle.close() + await monitor.automaticRefresh() + XCTAssertEqual(monitor.readerSnapshot.localTokenActivity.tokens, 600) + + exploration.timeRange = .selected + exploration.visibleRange = DateInterval( + start: fetchedAt.addingTimeInterval(-3_600), + end: fetchedAt + ) + monitor.analyticsPreferencesDidChange( + exploration: exploration, + dispositions: [:] + ) + for _ in 0 ..< 200 + where monitor.readerSnapshot.localTokenActivity.interval.duration + < 80 * 86_400 { + try await Task.sleep(for: .milliseconds(10)) } - try await Task.sleep(for: .milliseconds(10)) - await gate.release() - await load.value - await hide.value + XCTAssertGreaterThanOrEqual( + monitor.readerSnapshot.localTokenActivity.interval.duration, + 80 * 86_400 + ) - XCTAssertNil(monitor.readerSnapshot.localTokenActivity.tokens) + exploration.filters.model = "gpt-5.6" + monitor.analyticsPreferencesDidChange( + exploration: exploration, + dispositions: [:] + ) + for _ in 0 ..< 200 + where monitor.readerSnapshot.localTokenActivity.tokens != 0 { + try await Task.sleep(for: .milliseconds(10)) + } + requestCount = await requests.count + XCTAssertEqual(requestCount, 0) + XCTAssertEqual(monitor.readerSnapshot.localTokenActivity.tokens, 0) } func testMonitorFinishesABoundedLocalImportAcrossAutomaticRefresh() @@ -1016,7 +1133,12 @@ final class UsageMonitorHistoryTests: XCTestCase { fetchUsage: { try await fetches.next() } ) - await monitor.setLocalAnalyticsVisible(true) + var exploration = AnalyticsExplorationState.initial + exploration.graph = .tokenActivity + await monitor.setLocalAnalyticsVisible( + true, + exploration: exploration + ) await monitor.refresh() try Data(rollout.utf8).write(to: rolloutURL) await monitor.refresh() @@ -2003,33 +2125,6 @@ private actor RequestCounter { var count: Int { requests } } -private actor ProjectionGate { - private var started = false - private var continuation: CheckedContinuation? - - func response() async -> Data { - started = true - return await withCheckedContinuation { continuation in - self.continuation = continuation - } - } - - func waitUntilStarted() async { - while !started { - await Task.yield() - } - } - - func release() { - continuation?.resume( - returning: Data( - #"{"result":{"data":[],"nextCursor":null}}"#.utf8 - ) - ) - continuation = nil - } -} - private actor DelayedFetchSource { private let result: CodexFetchResult private var calls = 0 @@ -2053,6 +2148,8 @@ private final class BlockingUsageEvaluator: @unchecked Sendable { private let started: XCTestExpectation private var blockedFirstAccountEvaluation = false private var inputs: [UsageIntelligenceInput] = [] + private var activeCalls = 0 + private var peakActiveCalls = 0 init(started: XCTestExpectation) { self.started = started @@ -2063,6 +2160,8 @@ private final class BlockingUsageEvaluator: @unchecked Sendable { ) -> UsageReaderSnapshot { lock.lock() inputs.append(input) + activeCalls += 1 + peakActiveCalls = max(peakActiveCalls, activeCalls) let shouldBlock = input.account != nil && !blockedFirstAccountEvaluation if shouldBlock { @@ -2074,7 +2173,11 @@ private final class BlockingUsageEvaluator: @unchecked Sendable { started.fulfill() gate.wait() } - return UsageIntelligenceEngine.evaluate(input) + let result = UsageIntelligenceEngine.evaluate(input) + lock.lock() + activeCalls -= 1 + lock.unlock() + return result } func release() { @@ -2092,6 +2195,12 @@ private final class BlockingUsageEvaluator: @unchecked Sendable { defer { lock.unlock() } return inputs.last?.analyticsExploration } + + var maximumConcurrentCalls: Int { + lock.lock() + defer { lock.unlock() } + return peakActiveCalls + } } private struct StoredStateFixture: Codable {