From d7b02477b210e02fff9d0bdea30b2201a91b1a02 Mon Sep 17 00:00:00 2001
From: thrr87 <193831865+thrr87@users.noreply.github.com>
Date: Wed, 29 Jul 2026 16:54:58 +0200
Subject: [PATCH] fix: bound usage chart and harden Codex reads

---
 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()
             }
         )