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 27f0538c18edb6c5f156413af352f0150ec7f5ce Mon Sep 17 00:00:00 2001 From: thrr87 Date: Wed, 29 Jul 2026 17:16:57 +0200 Subject: [PATCH 2/5] fix: bound usage chart and harden Codex reads (#34) Co-authored-by: thrr87 <193831865+thrr87@users.noreply.github.com> --- Sources/CodexLimits/CodexClient.swift | 27 ++++++++++--- Sources/CodexLimits/MenuContentView.swift | 2 +- .../CodexLimits/UsageIntelligenceEngine.swift | 14 +++++++ .../AnalyticsWorkspaceTests.swift | 39 +++++++++++++++++++ Tests/CodexLimitsTests/CodexClientTests.swift | 34 ++++++++++++++++ 5 files changed, 110 insertions(+), 6 deletions(-) diff --git a/Sources/CodexLimits/CodexClient.swift b/Sources/CodexLimits/CodexClient.swift index 5e96988..be3c0df 100644 --- a/Sources/CodexLimits/CodexClient.swift +++ b/Sources/CodexLimits/CodexClient.swift @@ -70,9 +70,9 @@ enum CodexClientError: LocalizedError { final class CodexAppServerConnection: @unchecked Sendable { let input: FileHandle - let output: FileHandle let isRunning: () -> Bool let stop: () -> Void + private let outputDescriptor: Int32 private var bufferedOutput = Data() init( @@ -82,9 +82,15 @@ final class CodexAppServerConnection: @unchecked Sendable { stop: @escaping () -> Void ) { self.input = input - self.output = output self.isRunning = isRunning self.stop = stop + outputDescriptor = Darwin.dup(output.fileDescriptor) + } + + deinit { + if outputDescriptor >= 0 { + Darwin.close(outputDescriptor) + } } func readLine() async -> Data? { @@ -94,10 +100,21 @@ final class CodexAppServerConnection: @unchecked Sendable { bufferedOutput.removeSubrange(...newline) return Data(line) } - let chunk = await Task.detached { [output] in - output.availableData + let chunk = await Task.detached { + [outputDescriptor] () -> Data? in + var data = Data(count: 64 * 1_024) + let count = data.withUnsafeMutableBytes { + Darwin.read( + outputDescriptor, + $0.baseAddress, + $0.count + ) + } + guard count > 0 else { return nil } + data.count = count + return data }.value - guard !chunk.isEmpty else { + guard let chunk else { guard !bufferedOutput.isEmpty else { return nil } defer { bufferedOutput.removeAll() } return bufferedOutput diff --git a/Sources/CodexLimits/MenuContentView.swift b/Sources/CodexLimits/MenuContentView.swift index 75c093f..9868520 100644 --- a/Sources/CodexLimits/MenuContentView.swift +++ b/Sources/CodexLimits/MenuContentView.swift @@ -2316,7 +2316,7 @@ private struct UsageRemainingChart: View { @ChartContentBuilder private var observedMarks: some ChartContent { ForEach( - Array(chart.allObservedSegments.enumerated()), + Array(chart.observedSegments(within: visibleRange).enumerated()), id: \.offset ) { segmentIndex, segment in ForEach(segment) { point in diff --git a/Sources/CodexLimits/UsageIntelligenceEngine.swift b/Sources/CodexLimits/UsageIntelligenceEngine.swift index 489fa84..ab85fd1 100644 --- a/Sources/CodexLimits/UsageIntelligenceEngine.swift +++ b/Sources/CodexLimits/UsageIntelligenceEngine.swift @@ -305,6 +305,20 @@ struct UsageChartSnapshot: Equatable, Sendable { allObservedSegments.flatMap { $0 } } + func observedSegments( + within range: DateInterval + ) -> [[UsageChartPoint]] { + allowanceWindows + .filter { $0.resetsAt > range.start } + .flatMap(\.observedSegments) + .compactMap { segment in + let visible = segment.filter { + $0.date >= range.start && $0.date <= range.end + } + return visible.isEmpty ? nil : visible + } + } + var historicalProjection: [UsageChartPoint] { reference?.source == .accountHistory ? reference?.points ?? [] : [] } diff --git a/Tests/CodexLimitsTests/AnalyticsWorkspaceTests.swift b/Tests/CodexLimitsTests/AnalyticsWorkspaceTests.swift index 15f9fb6..1b25518 100644 --- a/Tests/CodexLimitsTests/AnalyticsWorkspaceTests.swift +++ b/Tests/CodexLimitsTests/AnalyticsWorkspaceTests.swift @@ -400,6 +400,45 @@ final class AnalyticsWorkspaceTests: XCTestCase { ) } + func testObservedSegmentsWithinCurrentWindowExcludePriorWindow() { + let currentWindow = DateInterval( + start: Date(timeIntervalSince1970: 4_000), + end: Date(timeIntervalSince1970: 8_000) + ) + let currentPoint = UsageChartPoint( + date: Date(timeIntervalSince1970: 5_000), + remaining: 70 + ) + let chart = UsageChartSnapshot( + observedSource: .account, + target: [], + currentProjection: [], + currentAllowanceReset: currentWindow.end, + allowanceWindows: [ + UsageAllowanceWindowSeries( + resetsAt: currentWindow.start, + observedSegments: [[ + UsageChartPoint( + date: currentWindow.start, + remaining: 40 + ) + ]] + ), + UsageAllowanceWindowSeries( + resetsAt: currentWindow.end, + observedSegments: [[currentPoint]] + ) + ], + currentRunsFaster: false, + accessibilityValue: "Observed usage" + ) + + XCTAssertEqual( + chart.observedSegments(within: currentWindow), + [[currentPoint]] + ) + } + func testHistoricalUsagePresetsReachBeyondTheCurrentWindow() { let suite = "AnalyticsWorkspaceTests.historicalUsagePresets" let defaults = UserDefaults(suiteName: suite)! diff --git a/Tests/CodexLimitsTests/CodexClientTests.swift b/Tests/CodexLimitsTests/CodexClientTests.swift index 0c6ce0b..da428a8 100644 --- a/Tests/CodexLimitsTests/CodexClientTests.swift +++ b/Tests/CodexLimitsTests/CodexClientTests.swift @@ -112,6 +112,39 @@ final class CodexClientTests: XCTestCase { XCTAssertEqual(server.initializationCount, 1) } + func testClosedServerOutputReportsConnectionLost() async { + let client = CodexClient( + makeConnection: { + let requests = Pipe() + let responses = Pipe() + let connection = CodexAppServerConnection( + input: requests.fileHandleForWriting, + output: responses.fileHandleForReading, + isRunning: { true }, + stop: { + try? requests.fileHandleForWriting.close() + try? responses.fileHandleForWriting.close() + } + ) + try responses.fileHandleForReading.close() + try responses.fileHandleForWriting.close() + return connection + }, + timeout: 1 + ) + + do { + _ = try await client.fetch( + fetchedAt: Date(timeIntervalSince1970: 1_900_000) + ) + XCTFail("Expected the closed connection to fail") + } catch CodexClientError.connectionLost { + // Expected. + } catch { + XCTFail("Expected connectionLost, got \(error)") + } + } + func testThreadProjectionReadsReuseTheInitializedAccountSession() async throws { let server = PersistentAppServerFixture() let client = CodexClient( @@ -1288,6 +1321,7 @@ private final class PersistentAppServerFixture: @unchecked Sendable { isRunning: { true }, stop: { try? requests.fileHandleForWriting.close() + try? responses.fileHandleForReading.close() try? responses.fileHandleForWriting.close() } ) 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 3/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 4/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 83791ca4a83e7a4cf79f3ce2937e2a43d1c9a1d1 Mon Sep 17 00:00:00 2001 From: thrr87 Date: Thu, 30 Jul 2026 21:50:18 +0200 Subject: [PATCH 5/5] Load local analytics only when needed (#37) * perf: load local analytics on demand * Use App Server account data for the lightweight core (#42) * perf: use account data for lightweight analytics core * chore: prepare v0.2.4 --------- Co-authored-by: thrr87 <193831865+thrr87@users.noreply.github.com> --------- Co-authored-by: thrr87 <193831865+thrr87@users.noreply.github.com> --- Resources/Info.plist | 4 +- Sources/CodexLimits/AnalyticsWorkspace.swift | 74 +++- .../CodexLimits/LocalActivityCollector.swift | 19 +- Sources/CodexLimits/MenuContentView.swift | 362 +++++++----------- Sources/CodexLimits/UsageMonitor.swift | 139 +++++-- .../AnalyticsWorkspaceTests.swift | 136 ++++++- .../LocalActivityCollectorTests.swift | 107 ++++++ .../UsageMonitorHistoryTests.swift | 225 ++++++++++- docs/research/app-server-architecture.md | 245 ++++++++++++ 9 files changed, 1051 insertions(+), 260 deletions(-) create mode 100644 docs/research/app-server-architecture.md diff --git a/Resources/Info.plist b/Resources/Info.plist index dcc9964..ee35d81 100644 --- a/Resources/Info.plist +++ b/Resources/Info.plist @@ -13,9 +13,9 @@ CFBundlePackageType APPL CFBundleShortVersionString - 0.2.2 + 0.2.4 CFBundleVersion - 4 + 5 LSApplicationCategoryType public.app-category.developer-tools LSMinimumSystemVersion diff --git a/Sources/CodexLimits/AnalyticsWorkspace.swift b/Sources/CodexLimits/AnalyticsWorkspace.swift index 5afc9f5..eda7962 100644 --- a/Sources/CodexLimits/AnalyticsWorkspace.swift +++ b/Sources/CodexLimits/AnalyticsWorkspace.swift @@ -18,8 +18,19 @@ enum AnalyticsGraph: String, CaseIterable, Codable, Identifiable, Sendable { var id: String { rawValue } + static let coreCases: [AnalyticsGraph] = [ + .usageRemaining, + .tokenActivity + ] + var usesAccountScope: Bool { - self == .usageRemaining || self == .usagePerToken + self == .usageRemaining + || self == .tokenActivity + || self == .usagePerToken + } + + var usesLocalAnalytics: Bool { + self == .usagePerToken || self == .concurrency } } @@ -62,6 +73,57 @@ enum AnalyticsTimeRange: String, CaseIterable, Codable, Identifiable, Sendable { } } +struct AccountTokenActivityRange: Equatable, Sendable { + let days: [TokenDay] + let completeDayCount: Int + let completeTokens: Int64? + + init(days: [TokenDay], interval: DateInterval) { + let day: TimeInterval = 86_400 + self.days = days + .filter { + $0.date < interval.end + && $0.date.addingTimeInterval(day) > interval.start + } + .sorted { $0.date < $1.date } + + var calendar = Calendar(identifier: .gregorian) + calendar.timeZone = TimeZone(secondsFromGMT: 0) + ?? calendar.timeZone + let startOfDay = calendar.startOfDay(for: interval.start) + var expected = startOfDay < interval.start + ? startOfDay.addingTimeInterval(day) + : startOfDay + var expectedDates: [Date] = [] + while expected.addingTimeInterval(day) <= interval.end { + expectedDates.append(expected) + expected = expected.addingTimeInterval(day) + } + completeDayCount = expectedDates.count + + guard !expectedDates.isEmpty else { + completeTokens = nil + return + } + var total: Int64 = 0 + for date in expectedDates { + guard let bucket = self.days.first(where: { $0.date == date }), + bucket.completeness == .complete, + bucket.tokens >= 0 else { + completeTokens = nil + return + } + let result = total.addingReportingOverflow(bucket.tokens) + guard !result.overflow else { + completeTokens = nil + return + } + total = result.partialValue + } + completeTokens = total + } +} + struct WorkspaceFilters: Codable, Equatable, Sendable { var projectID: String? var taskTreeID: String? @@ -101,6 +163,10 @@ struct AnalyticsExplorationState: Codable, Equatable, Sendable { pinnedUsageBaselineID: nil, pinnedUsageBaselineAccountPartitionID: nil ) + + var usesLocalAnalytics: Bool { + section == .graphs && graph.usesLocalAnalytics + } } @MainActor @@ -133,6 +199,12 @@ final class AnalyticsWorkspaceStore: ObservableObject { ) else { return .initial } + guard AnalyticsGraph.coreCases.contains(restored.graph) else { + var core = restored + core.graph = .usageRemaining + core.filters = .all + return core + } return restored } 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..1749f10 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) { + let state = workspace.state + await monitor.setLocalAnalyticsVisible( + state.usesLocalAnalytics + ) + if state.section == .graphs, + state.graph == .tokenActivity { + await monitor.refreshAccountIfStale() + } + } .environment(\.locale, Locale(identifier: "en_US")) } @@ -486,7 +495,7 @@ private struct GraphsWorkspace: View { set: store.selectGraph ) ) { - ForEach(AnalyticsGraph.allCases) { graph in + ForEach(AnalyticsGraph.coreCases) { graph in Text(graph.rawValue).tag(graph) } } @@ -518,9 +527,7 @@ private struct GraphsWorkspace: View { Label("Account", systemImage: "person.crop.circle") .font(.caption) .foregroundStyle(.secondary) - .help( - "Project, Task Tree, model, and reasoning filters do not change \(store.state.graph.rawValue)." - ) + .help("Data from your Codex account.") .accessibilityLabel("Account scope") } else { WorkspaceFilterMenu(reader: reader, store: store) @@ -1511,66 +1518,72 @@ private struct TokenActivityWorkspace: View { let reader: UsageReaderSnapshot @ObservedObject var store: AnalyticsWorkspaceStore - @State private var selectedPoint: LocalTokenActivityPoint? + @State private var selectedDay: TokenDay? + private let day: TimeInterval = 86_400 + + private var accountDays: [TokenDay] { + (reader.account?.tokenHistory ?? []).sorted { $0.date < $1.date } + } + + private var currentWindowBounds: DateInterval? { + reader.weeklyUsageRemaining.map { + DateInterval( + start: $0.window.startsAt, + end: $0.window.resetsAt + ) + } + } private var bounds: DateInterval { - reader.localTokenActivity.interval + let fallback = reader.fetchedAt ?? Date() + let first = accountDays.first?.date + ?? currentWindowBounds?.start + ?? fallback.addingTimeInterval(-day) + let last = accountDays.last?.date.addingTimeInterval(day) + ?? currentWindowBounds?.end + ?? fallback + return DateInterval( + start: min(first, currentWindowBounds?.start ?? first), + end: max(last, currentWindowBounds?.end ?? last) + ) } private var visibleRange: DateInterval { - store.effectiveRange( + if store.state.timeRange == .currentWindow, + let currentWindowBounds { + return currentWindowBounds + } + return store.effectiveRange( within: bounds, endingAt: min( - reader.localTokenActivity.observedAt ?? bounds.end, + reader.fetchedAt ?? bounds.end, bounds.end ) ) } - 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 + private var accountRange: AccountTokenActivityRange { + AccountTokenActivityRange( + days: accountDays, + interval: visibleRange ) } - private var accountCoversVisibleRange: Bool { - guard let interval = reader.accountTokenActivity.interval else { - return false - } - return abs(interval.start.timeIntervalSince(visibleRange.start)) < 1 - && abs(interval.end.timeIntervalSince(visibleRange.end)) < 1 - } - var body: some View { - content(localSlice) + content(accountRange) } - private func content(_ slice: LocalTokenActivitySlice) -> some View { + private func content(_ range: AccountTokenActivityRange) -> some View { VStack(alignment: .leading, spacing: 16) { VStack(alignment: .leading, spacing: 4) { Text("Token activity") .font(.title3.weight(.semibold)) - Text( - "Account and local token counts may differ, so we show them separately." - ) + Text("Daily token totals from your Codex account.") .font(.callout) .foregroundStyle(.secondary) } - ViewThatFits(in: .horizontal) { - HStack(alignment: .top, spacing: 12) { - accountCard - localCard(slice) - } - VStack(spacing: 12) { - accountCard - localCard(slice) - } - } + accountCard(range) VStack(alignment: .leading, spacing: 10) { ViewThatFits(in: .horizontal) { @@ -1585,33 +1598,35 @@ private struct TokenActivityWorkspace: View { } } - if slice.points.isEmpty { + if range.days.isEmpty { WorkspaceMessage( icon: "chart.xyaxis.line", - title: "No local token activity", - message: localEmptyMessage(slice) + title: "No account token activity", + message: "Codex did not return daily totals for this range." ) { EmptyView() } .frame(minHeight: 170) } else { - localChart(slice) + accountChart(range) } - selectedPointDetail(slice) + selectedPointDetail(range) } } .onChange(of: visibleRange) { _, range in - if let selectedPoint, !range.contains(selectedPoint.date) { - self.selectedPoint = nil + if let selectedDay, + selectedDay.date >= range.end + || selectedDay.date.addingTimeInterval(day) <= range.start { + self.selectedDay = nil } } } private var chartSourceLabel: some View { ChartLegendItem( - label: "Local Codex records", - color: .purple + label: "Daily totals · Account", + color: .blue ) } @@ -1621,42 +1636,24 @@ private struct TokenActivityWorkspace: View { .foregroundStyle(.tertiary) } - private var accountCard: some View { + private func accountCard( + _ range: AccountTokenActivityRange + ) -> some View { TokenSourceCard( title: "Account", - source: accountSource, - value: accountValue, - detail: accountDetail, - coverage: accountCoverage, - freshness: reader.accountTokenActivity.interval?.end, - freshnessLabel: "Through", - color: .blue - ) - } - - private func localCard(_ slice: LocalTokenActivitySlice) -> some View { - TokenSourceCard( - title: "Local", - source: "Local Codex records", - value: reader.localTokenActivity.tokens == nil - ? "Not available" - : compactTokenCount(slice.tokens), - detail: localDetail(slice), - coverage: coverageName(slice.coverage), - freshness: reader.localTokenActivity.observedAt, + source: store.state.timeRange == .currentWindow + ? accountSource + : "Codex daily token totals", + value: summaryTokens(in: range).map(compactTokenCount) + ?? "Not available", + detail: summaryDetail(in: range), + coverage: summaryCoverage(in: range), + freshness: reader.fetchedAt, freshnessLabel: "Updated", - color: .purple + color: .blue ) } - private var accountValue: String { - guard accountCoversVisibleRange, - let tokens = reader.accountTokenActivity.tokens else { - return "Not available" - } - return compactTokenCount(tokens) - } - private var accountSource: String { switch reader.accountTokenActivity.method { case .lifetimeDelta: @@ -1668,13 +1665,25 @@ private struct TokenActivityWorkspace: View { } } - private var accountDetail: String { - guard accountCoversVisibleRange else { - if reader.accountTokenActivity.tokens != nil { - return "No account total for this selected range" + private func summaryTokens( + in range: AccountTokenActivityRange + ) -> Int64? { + if store.state.timeRange == .currentWindow { + return reader.accountTokenActivity.tokens + } + return range.completeTokens + } + + private func summaryDetail( + in range: AccountTokenActivityRange + ) -> String { + guard store.state.timeRange == .currentWindow else { + if range.completeDayCount == 0 { + return "No full days in this range" } - return reader.accountTokenActivity.reason - ?? "Account token activity is unavailable" + return range.completeTokens == nil + ? "Codex returned only part of this range" + : "Sum of \(range.completeDayCount) complete days" } switch reader.accountTokenActivity.method { case .lifetimeDelta: @@ -1689,8 +1698,14 @@ private struct TokenActivityWorkspace: View { } } - private var accountCoverage: String { - guard accountCoversVisibleRange else { return "Unavailable" } + private func summaryCoverage( + in range: AccountTokenActivityRange + ) -> String { + guard store.state.timeRange == .currentWindow else { + return range.completeTokens == nil + ? "Unavailable" + : "Complete days" + } switch reader.accountTokenActivity.state { case .exact: return "Complete" case .partial: return "Partial" @@ -1698,57 +1713,33 @@ private struct TokenActivityWorkspace: View { } } - private func localDetail(_ slice: LocalTokenActivitySlice) -> String { - var details: [String] = [] - if let version = reader.localTokenActivity.sourceVersion { - details.append("Codex \(version)") - } - if let reason = slice.reason { - details.append(readerFacingLocalReason(reason)) - } - return details.isEmpty - ? "Local Codex records are unavailable" - : details.joined(separator: " · ") - } - - private func localEmptyMessage(_ slice: LocalTokenActivitySlice) -> String { - slice.reason.map(readerFacingLocalReason) - ?? "No local token events were found in this range." - } - - 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)] - + downsampledForDisplay(points, limit: 999) + private func renderedDays( + _ range: AccountTokenActivityRange + ) -> [TokenDay] { + downsampledForDisplay(range.days) } - private func localChart(_ slice: LocalTokenActivitySlice) -> some View { + private func accountChart( + _ range: AccountTokenActivityRange + ) -> some View { Chart { - ForEach(renderedChartPoints(slice)) { point in - LineMark( - x: .value("Time", point.date), - y: .value("Local tokens", point.tokens), - series: .value("Source", "Local Codex records") + ForEach(renderedDays(range), id: \.date) { tokenDay in + BarMark( + x: .value("Day", tokenDay.date, unit: .day), + y: .value("Account tokens", tokenDay.tokens) ) - .foregroundStyle(Color.purple) - .lineStyle(StrokeStyle(lineWidth: 2)) - .interpolationMethod(.stepEnd) + .foregroundStyle(Color.blue) } - if let selectedPoint { - RuleMark(x: .value("Selected time", selectedPoint.date)) + if let selectedDay { + RuleMark(x: .value("Selected time", selectedDay.date)) .foregroundStyle(Color.primary.opacity(0.45)) .lineStyle(StrokeStyle(lineWidth: 1, dash: [3, 3])) PointMark( - x: .value("Selected time", selectedPoint.date), - y: .value("Local tokens", selectedPoint.tokens) + x: .value("Selected day", selectedDay.date), + y: .value("Account tokens", selectedDay.tokens) ) - .foregroundStyle(Color.purple) + .foregroundStyle(Color.blue) .symbolSize(52) } } @@ -1777,21 +1768,21 @@ private struct TokenActivityWorkspace: View { at: location, proxy: proxy, geometry: geometry, - points: slice.points + days: range.days ) case .ended: - selectedPoint = nil + selectedDay = nil } } } } - .frame(height: 220) + .frame(height: 180) .accessibilityElement(children: .ignore) - .accessibilityLabel("Local token activity") + .accessibilityLabel("Account token activity") .accessibilityValue( - selectedPoint.map { - "\(compactTokenCount($0.tokens)) local tokens, \($0.date.formatted(date: .abbreviated, time: .shortened))" - } ?? "\(compactTokenCount(slice.tokens)) local tokens in the selected range" + selectedDay.map { + "\(compactTokenCount($0.tokens)) account tokens, \($0.date.formatted(date: .abbreviated, time: .omitted))" + } ?? "Daily account token totals are shown." ) .accessibilityHint( "Use Previous point and Next point for exact values." @@ -1800,36 +1791,36 @@ private struct TokenActivityWorkspace: View { @ViewBuilder private func selectedPointDetail( - _ slice: LocalTokenActivitySlice + _ range: AccountTokenActivityRange ) -> some View { VStack(alignment: .leading, spacing: 7) { HStack(spacing: 10) { - if let selectedPoint { + if let selectedDay { ViewThatFits(in: .horizontal) { HStack(spacing: 10) { - Text("Local Codex records") + Text("Daily account total") .fontWeight(.semibold) - Text(compactTokenCount(selectedPoint.tokens)) + Text(compactTokenCount(selectedDay.tokens)) .monospacedDigit() Text( - selectedPoint.date.formatted( + selectedDay.date.formatted( date: .abbreviated, - time: .shortened + time: .omitted ) ) .foregroundStyle(.secondary) } VStack(alignment: .leading, spacing: 3) { HStack(spacing: 8) { - Text("Local Codex records") + Text("Daily account total") .fontWeight(.semibold) - Text(compactTokenCount(selectedPoint.tokens)) + Text(compactTokenCount(selectedDay.tokens)) .monospacedDigit() } Text( - selectedPoint.date.formatted( + selectedDay.date.formatted( date: .abbreviated, - time: .shortened + time: .omitted ) ) .foregroundStyle(.secondary) @@ -1841,28 +1832,18 @@ private struct TokenActivityWorkspace: View { } Spacer() Button { - moveSelection(in: slice.points, by: -1) + moveSelection(in: range.days, by: -1) } label: { Image(systemName: "chevron.left") } .accessibilityLabel("Previous point") Button { - moveSelection(in: slice.points, by: 1) + moveSelection(in: range.days, by: 1) } label: { Image(systemName: "chevron.right") } .accessibilityLabel("Next point") } - if selectedPoint != nil { - HStack(spacing: 10) { - Text("Account · \(accountSource)") - .fontWeight(.semibold) - Text(accountValue) - .monospacedDigit() - Text("selected range") - .foregroundStyle(.secondary) - } - } } .font(.caption) .padding(10) @@ -1874,80 +1855,34 @@ private struct TokenActivityWorkspace: View { } private func moveSelection( - in points: [LocalTokenActivityPoint], + in days: [TokenDay], by offset: Int ) { - selectedPoint = steppedPoint( - in: points, - from: selectedPoint, + selectedDay = steppedPoint( + in: days, + from: selectedDay, by: offset ) } - private func readerFacingLocalReason(_ reason: String) -> String { - switch reason { - case "Local token activity starts from an unbounded counter": - "The first local reading has no earlier reading" - case "Local rollout path is unavailable", - "Local task records are missing": - "Some local Codex records could not be found" - case "Local task discovery is incomplete", - "Local task metadata is incomplete", - "Local task identity is missing": - "Some local tasks could not be checked" - case "This Codex CLI version has not been checked": - "This Codex version has not been checked" - case "Installed Codex CLI version is unavailable", - "Codex CLI version is unavailable": - "The installed Codex version could not be checked" - case "Only local activity on this Mac is observed": - "Only activity on this Mac is included" - case "Saved local activity could not be read": - "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 - } - } - private func selectNearestPoint( at location: CGPoint, proxy: ChartProxy, geometry: GeometryProxy, - points: [LocalTokenActivityPoint] + days: [TokenDay] ) { guard let date = chartDate( at: location, proxy: proxy, geometry: geometry ) else { return } - selectedPoint = nearestPoint( - in: points, + selectedDay = nearestPoint( + in: days, to: date, date: \.date ) } - private func coverageName(_ coverage: CoverageLevel) -> String { - switch coverage { - case .complete: "Complete" - case .high: "High" - case .partial: "Partial" - case .low: "Low" - case .unavailable: "Unavailable" - case .notApplicable: "Not applicable" - } - } - private func intervalText(_ interval: DateInterval) -> String { let start = interval.start.formatted( date: .abbreviated, @@ -2221,7 +2156,7 @@ private struct UsageRemainingChart: View { .chartOverlay { proxy in chartOverlay(proxy: proxy) } - .frame(height: 300) + .frame(height: 240) .accessibilityElement(children: .ignore) .accessibilityLabel("Usage remaining") .accessibilityValue( @@ -2791,13 +2726,6 @@ private struct FactsWorkspace: View { } } - WorkspaceCard(title: "Active Time") { - activeTimeContent - } - - WorkspaceCard(title: "Usage Receipts") { - receiptContent - } } } diff --git a/Sources/CodexLimits/UsageMonitor.swift b/Sources/CodexLimits/UsageMonitor.swift index ec6458a..c9b5d6a 100644 --- a/Sources/CodexLimits/UsageMonitor.swift +++ b/Sources/CodexLimits/UsageMonitor.swift @@ -14,6 +14,7 @@ enum SafetyBufferPolicy { @MainActor final class UsageMonitor: ObservableObject { + private static let accountRefreshInterval: TimeInterval = 600 static let safetyBufferKey = "safetyBuffer" @Published private(set) var readerSnapshot = UsageIntelligenceEngine.evaluate( @@ -80,6 +81,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( @@ -196,13 +199,15 @@ final class UsageMonitor: ObservableObject { ) } - Timer.publish(every: 600, on: .main, in: .common) + Timer.publish( + every: Self.accountRefreshInterval, + on: .main, + in: .common + ) .autoconnect() .sink { [weak self] _ in Task { - @MainActor in await self?.refresh( - forceHistorySync: false - ) + @MainActor in await self?.automaticRefresh() } } .store(in: &cancellables) @@ -211,21 +216,40 @@ 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 refresh(forceHistorySync: Bool = true) async { + func automaticRefresh() async { + await refresh( + forceHistorySync: false, + includeLocalActivity: false + ) + } + + func refreshAccountIfStale(now: Date = Date()) async { + guard let fetchedAt = accountSnapshot?.fetchedAt, + now.timeIntervalSince(fetchedAt) + < Self.accountRefreshInterval else { + await automaticRefresh() + return + } + } + + 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 +263,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 +321,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 +341,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 +355,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 +491,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 +538,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 +971,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 +989,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 +998,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..016326b 100644 --- a/Tests/CodexLimitsTests/AnalyticsWorkspaceTests.swift +++ b/Tests/CodexLimitsTests/AnalyticsWorkspaceTests.swift @@ -27,7 +27,7 @@ final class AnalyticsWorkspaceTests: XCTestCase { let first = AnalyticsWorkspaceStore(defaults: defaults) first.selectSection(.facts) - first.selectGraph(.concurrency) + first.selectGraph(.tokenActivity) first.selectTimeRange(.threeDays) first.updateFilters( WorkspaceFilters( @@ -51,7 +51,7 @@ final class AnalyticsWorkspaceTests: XCTestCase { let restored = AnalyticsWorkspaceStore(defaults: defaults) XCTAssertEqual(restored.state.section, .facts) - XCTAssertEqual(restored.state.graph, .concurrency) + XCTAssertEqual(restored.state.graph, .tokenActivity) XCTAssertEqual(restored.state.timeRange, .selected) XCTAssertEqual(restored.state.filters.projectID, "codex-limits") XCTAssertEqual(restored.state.filters.taskTreeID, "task-42") @@ -112,10 +112,75 @@ final class AnalyticsWorkspaceTests: XCTestCase { func testUsagePerTokenKeepsAccountScope() { XCTAssertTrue(AnalyticsGraph.usagePerToken.usesAccountScope) - XCTAssertFalse(AnalyticsGraph.tokenActivity.usesAccountScope) + XCTAssertTrue(AnalyticsGraph.tokenActivity.usesAccountScope) XCTAssertFalse(AnalyticsGraph.concurrency.usesAccountScope) } + func testLightweightCoreOffersOnlyAccountGraphs() { + XCTAssertEqual( + AnalyticsGraph.coreCases, + [.usageRemaining, .tokenActivity] + ) + } + + func testAccountTokenRangeSumsOnlyCompleteFullDays() throws { + let formatter = ISO8601DateFormatter() + let interval = DateInterval( + start: try XCTUnwrap( + formatter.date(from: "2026-07-01T12:00:00Z") + ), + end: try XCTUnwrap( + formatter.date(from: "2026-07-04T12:00:00Z") + ) + ) + let days = try [ + ("2026-07-01T00:00:00Z", 100, TokenDayCompleteness.complete), + ("2026-07-02T00:00:00Z", 200, .complete), + ("2026-07-03T00:00:00Z", 300, .complete), + ("2026-07-04T00:00:00Z", 400, .partial) + ].map { + TokenDay( + date: try XCTUnwrap(formatter.date(from: $0.0)), + tokens: Int64($0.1), + completeness: $0.2 + ) + } + + let range = AccountTokenActivityRange( + days: days, + interval: interval + ) + + XCTAssertEqual(range.days.count, 4) + XCTAssertEqual(range.completeDayCount, 2) + XCTAssertEqual(range.completeTokens, 500) + + let missingDay = AccountTokenActivityRange( + days: days.filter { $0.date != days[2].date }, + interval: interval + ) + XCTAssertNil(missingDay.completeTokens) + } + + func testRestoredLocalGraphFallsBackToUsageRemaining() throws { + let defaults = try XCTUnwrap( + UserDefaults( + suiteName: "AnalyticsWorkspaceTests-\(UUID().uuidString)" + ) + ) + var state = AnalyticsExplorationState.initial + state.graph = .concurrency + defaults.set( + try JSONEncoder().encode(state), + forKey: AnalyticsWorkspaceStore.persistenceKey + ) + + XCTAssertEqual( + AnalyticsWorkspaceStore.restoredState(from: defaults).graph, + .usageRemaining + ) + } + func testChangingGraphKeepsRangeAndFilters() { let store = AnalyticsWorkspaceStore( defaults: UserDefaults(suiteName: "AnalyticsWorkspaceTests-\(UUID().uuidString)")! @@ -146,7 +211,70 @@ final class AnalyticsWorkspaceTests: XCTestCase { func testUsageRemainingAlwaysUsesAccountScope() { XCTAssertTrue(AnalyticsGraph.usageRemaining.usesAccountScope) - XCTAssertFalse(AnalyticsGraph.tokenActivity.usesAccountScope) + XCTAssertTrue(AnalyticsGraph.tokenActivity.usesAccountScope) + } + + func testTokenActivityRendersAccountDailyBucketsWithoutLocalFacts() { + let fetchedAt = Date(timeIntervalSince1970: 10 * 86_400) + let account = UsageSnapshot( + mainLimit: LimitReading( + limitId: "weekly", + name: "Weekly", + window: UsageWindow( + remainingPercent: 75, + resetsAt: fetchedAt.addingTimeInterval(3 * 86_400), + durationMinutes: 10_080 + ) + ), + otherLimits: [], + tokenHistory: [ + TokenDay( + date: fetchedAt.addingTimeInterval(-2 * 86_400), + tokens: 1_000, + completeness: .complete + ), + TokenDay( + date: fetchedAt.addingTimeInterval(-86_400), + tokens: 2_000, + completeness: .complete + ), + TokenDay( + date: fetchedAt, + tokens: 500, + completeness: .partial + ) + ], + emergencyResetCount: 0, + fetchedAt: fetchedAt + ) + let reader = UsageIntelligenceEngine.evaluate( + UsageIntelligenceInput( + account: account, + samples: [], + safetyBuffer: 3, + sourceState: .available, + now: fetchedAt, + previousStatus: nil + ) + ) + let store = AnalyticsWorkspaceStore( + defaults: UserDefaults( + suiteName: "AnalyticsWorkspaceTests-\(UUID().uuidString)" + )! + ) + store.selectGraph(.tokenActivity) + + XCTAssertTrue( + renders( + AnalyticsWorkspaceBody( + reader: reader, + store: store, + assistedInsights: CodexAssistedInsightStore() + ), + size: CGSize(width: 640, height: 780) + ) + ) + XCTAssertNil(reader.localTokenActivity.tokens) } func testPresetRangeEndsAtLatestObservedTimeAndIsClampedToWindow() { 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..bb79e6f 100644 --- a/Tests/CodexLimitsTests/UsageMonitorHistoryTests.swift +++ b/Tests/CodexLimitsTests/UsageMonitorHistoryTests.swift @@ -587,6 +587,44 @@ final class UsageMonitorHistoryTests: XCTestCase { ) } + func testTokenActivityRefreshesOnlyWhenAccountDataIsStale() async throws { + let suiteName = "UsageMonitorHistoryTests-\(UUID().uuidString)" + let defaults = try XCTUnwrap(UserDefaults(suiteName: suiteName)) + defer { defaults.removePersistentDomain(forName: suiteName) } + let fetchedAt = Date(timeIntervalSince1970: 1_700_000) + let source = FetchSequence([ + makeFetchResult( + identity: "user@example.com", + fetchedAt: fetchedAt, + remaining: 90 + ), + makeFetchResult( + identity: "user@example.com", + fetchedAt: fetchedAt.addingTimeInterval(600), + remaining: 89 + ) + ]) + let monitor = UsageMonitor( + defaults: defaults, + historyDirectory: temporaryDirectory(), + startsAutomatically: false, + fetchUsage: { try await source.next() } + ) + + await monitor.refresh() + await monitor.refreshAccountIfStale( + now: fetchedAt.addingTimeInterval(599) + ) + var callCount = await source.callCount + XCTAssertEqual(callCount, 1) + + await monitor.refreshAccountIfStale( + now: fetchedAt.addingTimeInterval(600) + ) + callCount = await source.callCount + XCTAssertEqual(callCount, 2) + } + func testBlockedEvaluationDoesNotFreezeMainActorOrPublishAStaleSnapshot() async throws { let suiteName = "UsageMonitorHistoryTests-\(UUID().uuidString)" let defaults = try XCTUnwrap(UserDefaults(suiteName: suiteName)) @@ -802,6 +840,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 +860,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 +872,139 @@ 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) + + await monitor.refresh() + requestCount = await requests.count + XCTAssertEqual(requestCount, 0) + + for graph in AnalyticsGraph.coreCases { + var state = AnalyticsExplorationState.initial + state.graph = graph + await monitor.setLocalAnalyticsVisible( + state.usesLocalAnalytics + ) + } + 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 +1053,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 +1068,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 +1088,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 +2045,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 diff --git a/docs/research/app-server-architecture.md b/docs/research/app-server-architecture.md new file mode 100644 index 0000000..1c7e2f4 --- /dev/null +++ b/docs/research/app-server-architecture.md @@ -0,0 +1,245 @@ +# App Server as the light core + +Date: 2026-07-30 +Source baseline: installed stable `codex-cli 0.145.0`; official tag [`rust-v0.145.0`](https://github.com/openai/codex/tree/25af12f7e61572b0bc18ddb1008be543b91519b0), commit `25af12f7e61572b0bc18ddb1008be543b91519b0` + +## Decision + +Use Codex App Server as the only source for the default account view. + +The light core should: + +1. Reuse the existing persistent App Server process. +2. Read `account/rateLimits/read` and `account/usage/read` after connection. +3. Reconcile sparse rate-limit updates seen during account reads. +4. Keep a small local cache of the last good account responses. +5. Never scan rollout JSONL, import local token facts, or sync derived token facts in the default mode. + +This is a high-confidence cut for Codex CLI 0.145.0. + +The proposed `synced_token_fact` design may suit a later, optional local analytics mode. It should not enter the light core. The main Token Activity data already comes from the account, so a second cross-Mac copy would add work without adding data. + +## What the stable API supplies + +| Need | Stable source | What it supplies | Important limit | +|---|---|---|---| +| Usage remaining and reset time | `account/rateLimits/read` | Primary and secondary windows, `usedPercent`, window length, reset time, all named limit buckets | The percent is not a token quota | +| Live limit change | `account/rateLimits/updated` | One sparse rate-limit snapshot | It does not include all buckets or banked reset details | +| Banked resets | `account/rateLimits/read` | Authoritative available count and, when present, reset details and expiry dates | Detail rows may be absent or capped; there is no reset-credit event | +| Account Token Activity | `account/usage/read` | Daily token buckets and account summary facts | Daily only; no update event; retention and bucket time zone are not promised | +| Account facts | `account/usage/read` | Lifetime tokens, peak daily tokens, longest turn, current streak, longest streak | Each field may be absent | +| Local Task and project list | `thread/list` with `useStateDbOnly: true` | Stored Task metadata, working directory, parent link, source, dates and status | No token history, actual model ID, or reasoning effort | +| Live token use for one Task | `thread/tokenUsage/updated` | Thread ID, turn ID, cumulative and last token use, context-window size | Only for a Task started, forked, or resumed by that App Server connection | + +The contracts appear in the official [App Server guide](https://developers.openai.com/codex/app-server), [`account.rs`](https://github.com/openai/codex/blob/25af12f7e61572b0bc18ddb1008be543b91519b0/codex-rs/app-server-protocol/src/protocol/v2/account.rs), [`thread.rs`](https://github.com/openai/codex/blob/25af12f7e61572b0bc18ddb1008be543b91519b0/codex-rs/app-server-protocol/src/protocol/v2/thread.rs), and [`thread_data.rs`](https://github.com/openai/codex/blob/25af12f7e61572b0bc18ddb1008be543b91519b0/codex-rs/app-server-protocol/src/protocol/v2/thread_data.rs). They also appear in the stable schema generated by the installed CLI. + +None of these methods needs the experimental API flag in 0.145.0. + +## Account Token Activity needs no cross-Mac sync + +`account/usage/read` returns account data, not a scan of the current Mac. Two Macs signed in to the same Codex account and workspace should therefore receive the same daily Token Activity, subject to refresh time and backend delay. + +Codex Limits should not publish or import its own copy of those daily buckets. Each Mac can cache the last response for offline display. The cache is disposable and does not need folder sync. + +This removes the main reasons for the proposed remote-fact path: + +- no shared writer; +- no project projection merge; +- no remote generation import; +- no remote fact deduplication; +- no cross-Mac Token Activity coverage calculation. + +The same conclusion does **not** apply to future local Task receipts. Those facts describe work observed by one Mac. If the product later syncs them, the per-device manifest and separate `synced_token_fact` table are a sound starting boundary. That later design still needs a proven event identity. `thread/tokenUsage/updated` gives a cumulative snapshot, not a documented stable event ID, so deduplication cannot yet rest on a generic `eventID`. + +## There is no `account/usage/updated` + +The stable 0.145.0 protocol defines `account/usage/read`, but no account-usage notification. The App Server guide, generated stable schema, and official protocol source contain no `account/usage/updated` method. + +Token Activity therefore needs a read: + +- after the first connection; +- after reconnect; +- after wake when the cached response is old; +- on explicit refresh; +- on a slow schedule, or near a daily bucket boundary. + +Opening the menu should show the cache. It should not trigger a full read on every open. + +Rate limits differ. `account/rateLimits/updated` exists, but it is a sparse rolling update. The client must merge values into the last full snapshot by `limitId`. A missing optional value does not clear an older value. The event carries no banked reset summary, so Codex Limits must still read the full snapshot after reset use, reconnect, wake, or another reason to suspect missed events. The official source describes this merge-or-refetch rule in [`AccountRateLimitsUpdatedNotification`](https://github.com/openai/codex/blob/25af12f7e61572b0bc18ddb1008be543b91519b0/codex-rs/app-server-protocol/src/protocol/v2/account.rs). + +## One connection, not one process per refresh + +The official transport is JSONL over stdio. A client sends `initialize` exactly once for each transport connection, then `initialized`, and keeps reading responses and events. The App Server uses bounded queues; when it reports `-32001` because it is busy, the client should retry with exponential backoff and jitter. See the official [App Server README](https://github.com/openai/codex/blob/25af12f7e61572b0bc18ddb1008be543b91519b0/codex-rs/app-server/README.md). + +Codex Limits already has most of the process lifetime needed for this design: + +- [`CodexClient.shared`](../../Sources/CodexLimits/CodexClient.swift) owns a stored connection; +- `activeConnection()` reuses the running process while the Codex executable identity stays the same; +- `invalidateConnection()` stops it after a failed connection and allows a new one. + +The pivot does not need a second service or a new language. It needs a better reader around the existing connection. + +Today, Codex Limits reads stdout only while waiting for a request response. A true event-driven client would need one background reader that owns stdout for the life of the connection. That is a later transport change, not part of this lightweight-core cut. It should: + +1. Route responses to the pending request by ID. +2. Apply rate-limit events to the cache. +3. Treat EOF or a dead child as a lost connection. +4. Start a new process, initialize once, and read fresh account snapshots. +5. Back off after repeated failure. + +The official protocol does not promise replay of missed account events. Refetching both account snapshots after reconnect is therefore a client recovery rule, not a server guarantee. + +## `thread/list` avoids JSONL repair, but it is not account usage + +`thread/list` with `useStateDbOnly: true` reads the Codex state database without scanning rollout JSONL to repair metadata. Omit the flag and the server may scan rollouts to repair the list. + +This method is enough for a lazy local list of: + +- Tasks; +- Codex project labels derived from `cwd`; +- parent links already stored by Codex; +- created, updated, and recent dates. + +It is not enough for per-Task token use, model, reasoning, concurrency, or receipts. The stable `Thread` record includes `modelProvider`, but not the effective model ID or reasoning effort. + +The response also contains fields Codex Limits should not retain, such as the Task preview, full path, and Git metadata. The app should extract its small allowlist and drop the response. + +`useStateDbOnly: true` trades repair for speed. If Codex has not yet placed a Task in its state database, this call does not scan JSONL to recover it. That is the correct trade for the light core. + +## Task token events are not a global feed + +`thread/tokenUsage/updated` is an active-Task notification. A connection receives Task events after it starts, forks, or resumes that Task. `thread/list` and `thread/read` do not subscribe to the Task. + +The server can replay the last persisted cumulative token count after `thread/resume`, as shown by the official [`token_usage_replay.rs`](https://github.com/openai/codex/blob/25af12f7e61572b0bc18ddb1008be543b91519b0/codex-rs/app-server/src/request_processors/token_usage_replay.rs). This is not a history API: + +- it does not list past token events; +- it does not backfill every Task; +- it can rebuild persisted thread history; +- resuming a Task makes the caller a live subscriber and crosses the product's read-only Task boundary. + +A separate Codex Limits App Server process does not receive live Task events owned by another Codex process. Its subscription registry is process-local. This is covered in the prior [local activity source spike](local-activity-source-spike-2026-07-27.md). + +For the light core, do not use Task token events. For later optional analytics, collect them only for Tasks the analytics connection truly observes, and state the coverage. Do not promise complete historical backfill. + +## What `account/usage/read` can and cannot support + +It can support: + +- the main daily Token Activity graph; +- daily, weekly, four-week, and twelve-week sums when the returned buckets cover the selected range; +- lifetime and peak daily token facts; +- account trends shared across the user's Macs. + +It cannot support: + +- an exact hidden token allowance; +- a model, project, Task, agent, or reasoning breakdown; +- per-turn receipts; +- concurrency; +- exact intra-day timing; +- history older than the backend returns. + +Daily buckets also make partial boundary days uncertain. A four-week or twelve-week comparison should use complete returned days and show that limit. If the backend does not return the full range, the UI should say that the range is unavailable. It should not start a JSONL scan as a silent fallback. + +Account token activity and limit use are different measures. The app may show both, but it must not infer a true token quota from daily tokens and `usedPercent`. + +`account/usage/read` also requires Codex-service-backed authentication. API-key-only and Bedrock modes do not supply this account profile. In those modes, show Token Activity as unavailable. Do not turn on local history scanning without the user's action. + +## OpenTelemetry belongs to optional analytics + +Codex can export OpenTelemetry logs, traces, and metrics. The official [observability guide](https://developers.openai.com/codex/config-advanced#observability-and-telemetry) documents: + +- async OTLP export over HTTP or gRPC; +- per-turn token-use metrics split into input, cached input, output, and reasoning output; +- model and session tags; +- turn and tool timing; +- prompt logging off by default. + +It is a useful future source, but not a light-core dependency: + +- the user must change Codex config; +- Codex Limits must run or connect to an OTLP receiver; +- collection starts after enablement and has no promised history backfill; +- documented metrics are aggregates, not Task receipts; +- logs can contain tool-result snippets even when prompt content stays off. + +Use OpenTelemetry only behind a clear extra-analytics switch and a privacy review. Do not enable it or edit Codex config in the default mode. + +## Delivered light-core architecture + +```text +Codex App Server, one supervised stdio child + ├─ account/rateLimits/read ─┐ + ├─ rateLimits/updated ──────┼─> bounded account cache ─> Usage remaining + │ (seen during reads) │ + └─ account/usage/read ──────┘ └─> Token Activity + +Optional later analytics + ├─ thread/list(useStateDbOnly: true) + ├─ observed Task events + ├─ explicit, slow JSONL history import + └─ user-enabled OpenTelemetry +``` + +SQLite may remain as the bounded cache or local store. It is not the cause of the current load. The costly work comes from scanning, decoding, normalizing, projecting, and syncing data that the account API already supplies. + +## Remove or defer + +Remove from the light-core path: + +- rollout JSONL discovery and scan; +- automatic local token-fact import; +- task projections built for the account graph; +- cross-Mac Token Activity export and import; +- generation merge logic for account token activity; +- refresh on each menu open. + +Defer: + +- per-Task and per-agent token use; +- model and reasoning breakdown; +- concurrency; +- receipts; +- passive and Codex-assisted insights based on local Source Content; +- OTLP receiver; +- remote `synced_token_fact` and device manifests. + +The UI can keep these views hidden until their source is enabled and has data. It should not show a loading state for work the default mode never starts. + +## Delivery checks + +Before release, verify on both test Macs: + +1. Default launch and menu open read no rollout JSONL. +2. One Codex Limits process owns at most one ordinary App Server child. +3. Opening and closing the menu does not start another child or rescan history. +4. A sparse rate-limit event seen during a read triggers a bounded full + account reconciliation; opening the menu alone does not refresh. +5. Reconnect reads fresh rate limits, banked resets, account usage, and account identity. +6. Token Activity comes only from `account/usage/read`. +7. Two Macs on the same account and workspace show the same complete daily buckets within normal backend and refresh delay. +8. A missing or unsupported `account/usage/read` shows a clear unavailable state and starts no fallback scan. +9. Memory and refresh time stay within the agreed budgets on real account data. + +## First-machine measurement + +Measured on the developer Mac with real account data on 2026-07-30: + +- one Codex Limits process and one App Server child; +- no open rollout JSONL or local analytics store; +- steady `top` memory: about 45 MB for Codex Limits and 41–58 MB for App Server; +- peak refresh CPU: 19.2% for Codex Limits and 1.4% for App Server; +- no measured memory increase during refresh; +- Token Activity rendered in about one second, including the UI test tool's settle time. + +The second-Mac check remains a release gate. + +## Sources + +- [Codex App Server guide](https://developers.openai.com/codex/app-server) +- [App Server README and protocol lifecycle](https://github.com/openai/codex/blob/25af12f7e61572b0bc18ddb1008be543b91519b0/codex-rs/app-server/README.md) +- [Stable account protocol](https://github.com/openai/codex/blob/25af12f7e61572b0bc18ddb1008be543b91519b0/codex-rs/app-server-protocol/src/protocol/v2/account.rs) +- [Stable thread request protocol](https://github.com/openai/codex/blob/25af12f7e61572b0bc18ddb1008be543b91519b0/codex-rs/app-server-protocol/src/protocol/v2/thread.rs) +- [Stable Thread data](https://github.com/openai/codex/blob/25af12f7e61572b0bc18ddb1008be543b91519b0/codex-rs/app-server-protocol/src/protocol/v2/thread_data.rs) +- [Persisted token-usage replay](https://github.com/openai/codex/blob/25af12f7e61572b0bc18ddb1008be543b91519b0/codex-rs/app-server/src/request_processors/token_usage_replay.rs) +- [Observability and telemetry](https://developers.openai.com/codex/config-advanced#observability-and-telemetry)