Add cross-provider token usage accounting for agent transcripts

Packages/Shared/CmuxAgentChat/Sources/CmuxAgentChat/Model/ChatTokenUsage.swift

@@ -0,0 +1,125 @@
+import Foundation
+
+/// Token counts for one or more agent API responses, normalized across
+/// providers.
+///
+/// The providers disagree about what "input tokens" means, and the
+/// disagreement is silent: both write a key called `input_tokens` and the
+/// two keys mean different things.
+///
+/// - Claude Code's `input_tokens` counts *only* tokens that were neither
+///   read from nor written to the prompt cache. `cache_read_input_tokens`
+///   and `cache_creation_input_tokens` sit beside it and are not included
+///   in it.
+/// - Codex's `input_tokens` counts the *whole* prompt, with
+///   `cached_input_tokens` as a subset of it. Its own `total_tokens` is
+///   `input_tokens + output_tokens`, which only adds up because the cached
+///   part is already inside `input_tokens`.
+///
+/// So mapping both providers' `input_tokens` onto one field overstates
+/// Codex's uncached input by the entire cached prompt, which for a long
+/// session is nearly the whole thing. This type stores the three input
+/// kinds separately instead, and each extractor is responsible for
+/// converting into them.
+public struct ChatTokenUsage: Sendable, Equatable {
+    /// Input tokens that were neither served from nor written to the cache.
+    public var freshInputTokens: Int
+
+    /// Input tokens served from the prompt cache.
+    public var cacheReadTokens: Int
+
+    /// Input tokens written into the prompt cache.
+    public var cacheWriteTokens: Int
+
+    /// Tokens the model generated, including any reasoning tokens.
+    public var outputTokens: Int
+
+    /// The reasoning share of ``outputTokens``.
+    ///
+    /// Both providers bill reasoning as output and count it inside their
+    /// output total, so this is a breakdown and never an addend. Adding it
+    /// to a total double-counts it.
+    public var reasoningOutputTokens: Int
+
+    /// Creates a usage value.
+    ///
+    /// - Parameters:
+    ///   - freshInputTokens: Uncached input tokens.
+    ///   - cacheReadTokens: Input tokens read from the cache.
+    ///   - cacheWriteTokens: Input tokens written to the cache.
+    ///   - outputTokens: Generated tokens, reasoning included.
+    ///   - reasoningOutputTokens: The reasoning share of the output.
+    public init(
+        freshInputTokens: Int = 0,
+        cacheReadTokens: Int = 0,
+        cacheWriteTokens: Int = 0,
+        outputTokens: Int = 0,
+        reasoningOutputTokens: Int = 0
+    ) {
+        self.freshInputTokens = freshInputTokens
+        self.cacheReadTokens = cacheReadTokens
+        self.cacheWriteTokens = cacheWriteTokens
+        self.outputTokens = outputTokens
+        self.reasoningOutputTokens = reasoningOutputTokens
+    }
+
+    /// Every input token, cached or not.
+    public var inputTokens: Int {
+        Self.saturatedSum(
+            Self.saturatedSum(freshInputTokens, cacheReadTokens),
+            cacheWriteTokens
+        )
+    }
+
+    /// Every token, input and output.
+    ///
+    /// Deliberately excludes ``reasoningOutputTokens``, which is already
+    /// inside ``outputTokens``.
+    public var totalTokens: Int {
+        Self.saturatedSum(inputTokens, outputTokens)
+    }
+
+    /// Whether every count is zero.
+    public var isEmpty: Bool {
+        totalTokens == 0 && reasoningOutputTokens == 0
+    }
+
+    /// Adds two usage values field by field.
+    ///
+    /// - Parameters:
+    ///   - lhs: The left value.
+    ///   - rhs: The right value.
+    /// - Returns: The field-wise sum.
+    public static func + (lhs: ChatTokenUsage, rhs: ChatTokenUsage) -> ChatTokenUsage {
+        ChatTokenUsage(
+            freshInputTokens: Self.saturatedSum(lhs.freshInputTokens, rhs.freshInputTokens),
+            cacheReadTokens: Self.saturatedSum(lhs.cacheReadTokens, rhs.cacheReadTokens),
+            cacheWriteTokens: Self.saturatedSum(lhs.cacheWriteTokens, rhs.cacheWriteTokens),
+            outputTokens: Self.saturatedSum(lhs.outputTokens, rhs.outputTokens),
+            reasoningOutputTokens: Self.saturatedSum(
+                lhs.reasoningOutputTokens,
+                rhs.reasoningOutputTokens
+            )
+        )
+    }
+
+    /// Adds a usage value into this one.
+    ///
+    /// - Parameters:
+    ///   - lhs: The value to add into.
+    ///   - rhs: The value to add.
+    public static func += (lhs: inout ChatTokenUsage, rhs: ChatTokenUsage) {
+        lhs = lhs + rhs
+    }
+
+    /// Adds two integer counts without allowing extreme input to trap.
+    ///
+    /// A transcript's counts come from a provider, so a corrupt or hostile line
+    /// can carry a value near `Int.max`. Saturating keeps one bad line from
+    /// taking down the accumulation of every good one.
+    static func saturatedSum(_ lhs: Int, _ rhs: Int) -> Int {
+        let (sum, overflow) = lhs.addingReportingOverflow(rhs)
+        guard overflow else { return sum }
+        return lhs >= 0 ? Int.max : Int.min
+    }
+}

Packages/Shared/CmuxAgentChat/Sources/CmuxAgentChat/Model/ChatUsageCodexSource.swift

@@ -0,0 +1,15 @@
+/// Which Codex record type a transcript's accounting came from.
+///
+/// Once usage records appear they replace the cumulative fallback, whose
+/// lifetime total may have been inherited from a parent thread.
+public enum ChatUsageCodexSource: Sendable, Equatable {
+    /// No Codex usage seen yet.
+    case none
+
+    /// Per-response `token_usage_record` lines.
+    case usageRecords
+
+    /// The cumulative `total_token_usage` from `token_count` events, used
+    /// when the transcript predates `token_usage_record`.
+    case cumulativeEvents
+}

Packages/Shared/CmuxAgentChat/Sources/CmuxAgentChat/Model/ChatUsageRateLimit+Window.swift

@@ -0,0 +1,27 @@
+import Foundation
+
+extension ChatUsageRateLimit {
+    /// One rolling allowance window.
+    public struct Window: Sendable, Equatable {
+        /// Percentage of the window's allowance used, 0 to 100.
+        public var usedPercent: Double
+
+        /// Length of the rolling window in minutes.
+        public var windowMinutes: Int?
+
+        /// When the window resets, when the provider says.
+        public var resetsAt: Date?
+
+        /// Creates a window reading.
+        ///
+        /// - Parameters:
+        ///   - usedPercent: Percentage of the allowance used.
+        ///   - windowMinutes: Window length in minutes.
+        ///   - resetsAt: When the window resets.
+        public init(usedPercent: Double, windowMinutes: Int? = nil, resetsAt: Date? = nil) {
+            self.usedPercent = usedPercent
+            self.windowMinutes = windowMinutes
+            self.resetsAt = resetsAt
+        }
+    }
+}

Packages/Shared/CmuxAgentChat/Sources/CmuxAgentChat/Model/ChatUsageRateLimit.swift

@@ -0,0 +1,105 @@
+import Foundation
+
+/// A provider snapshot of allowance windows and spend-control state.
+///
+/// Codex writes this into its rollout on every token count. Claude Code
+/// does not put limit state in its transcript at all, so this stays `nil`
+/// for Claude sessions; the absence is the honest answer rather than a
+/// zero. Fields inside a Codex snapshot are independently optional too.
+public struct ChatUsageRateLimit: Sendable, Equatable {
+    /// The short window, when the provider reports it.
+    public var primary: Window?
+
+    /// The long window, when the provider reports one.
+    ///
+    /// Codex sends two windows: a roughly five-hour `primary` and a weekly
+    /// `secondary`. The weekly one is usually what actually stops a day of
+    /// work, and it is the one a session hits without warning, so dropping
+    /// it would hide the limit that matters.
+    public var secondary: Window?
+
+    /// Whether the provider says a spend control has cut the session off.
+    ///
+    /// `nil` means the snapshot omitted the flag; it must not be presented
+    /// as `false`, which would claim the control has not been reached.
+    public var spendControlReached: Bool?
+
+    /// Percentage of the ``primary`` window's allowance used.
+    public var usedPercent: Double? {
+        get { primary?.usedPercent }
+        set {
+            guard let newValue else {
+                primary = nil
+                return
+            }
+            if primary == nil {
+                primary = Window(usedPercent: newValue)
+            } else {
+                primary?.usedPercent = newValue
+            }
+        }
+    }
+
+    /// Length of the ``primary`` window in minutes.
+    public var windowMinutes: Int? {
+        get { primary?.windowMinutes }
+        set { primary?.windowMinutes = newValue }
+    }
+
+    /// When the ``primary`` window resets.
+    public var resetsAt: Date? {
+        get { primary?.resetsAt }
+        set { primary?.resetsAt = newValue }
+    }
+
+    /// Whichever reported window is closest to its limit.
+    ///
+    /// This is what a caller showing one number should show. Picking the
+    /// primary window instead reads "12% used" on a session that is at 96%
+    /// of its weekly allowance.
+    public var tightestWindow: Window? {
+        switch (primary, secondary) {
+        case let (primary?, secondary?):
+            return secondary.usedPercent > primary.usedPercent ? secondary : primary
+        case let (primary?, nil):
+            return primary
+        case let (nil, secondary?):
+            return secondary
+        case (nil, nil):
+            return nil
+        }
+    }
+
+    /// Creates a rate limit reading from its windows.
+    ///
+    /// - Parameters:
+    ///   - primary: The short window, when reported.
+    ///   - secondary: The long window, when reported.
+    ///   - spendControlReached: Whether a spend control has cut the session off,
+    ///     or `nil` when the snapshot omitted that state.
+    public init(
+        primary: Window? = nil,
+        secondary: Window? = nil,
+        spendControlReached: Bool? = nil
+    ) {
+        self.primary = primary
+        self.secondary = secondary
+        self.spendControlReached = spendControlReached
+    }
+
+    /// Creates a rate limit reading with only a primary window.
+    ///
+    /// - Parameters:
+    ///   - usedPercent: Percentage of the allowance used.
+    ///   - windowMinutes: Window length in minutes.
+    ///   - resetsAt: When the window resets.
+    public init(usedPercent: Double, windowMinutes: Int? = nil, resetsAt: Date? = nil) {
+        self.init(
+            primary: Window(
+                usedPercent: usedPercent,
+                windowMinutes: windowMinutes,
+                resetsAt: resetsAt
+            )
+        )
+    }
+}

Packages/Shared/CmuxAgentChat/Sources/CmuxAgentChat/Model/ChatUsageTotals.swift

@@ -0,0 +1,115 @@
+import Foundation
+
+/// Everything one transcript says about what it spent.
+///
+/// Carries no prices. Converting tokens to money needs a per-model price
+/// table, an answer for subscription plans where per-token prices do not
+/// apply, and a decision about how stale a bundled table may get. Those
+/// are product decisions, so this type stops at counts and leaves them to
+/// the caller.
+public struct ChatUsageTotals: Sendable, Equatable {
+    /// Usage summed over every distinct API response in the transcript.
+    public var usage: ChatTokenUsage
+
+    /// Usage split by the model that produced it.
+    ///
+    /// Keyed by the provider's own model identifier. A session that switched
+    /// models mid-run has an entry per model. A cumulative Codex fallback has
+    /// no response-level model identity, so the split can sum to less than
+    /// ``usage``.
+    public var usageByModel: [String: ChatTokenUsage]
+
+    /// Distinct API responses counted.
+    public var responses: Int
+
+    /// Repeated reports of a response that was already counted.
+    ///
+    /// Non-zero is normal, not a warning: both providers repeat usage by
+    /// design. A caller that wants to prove deduplication is working can
+    /// watch this climb, within one limit: identities are remembered in a
+    /// bounded recent window, so a transcript with more distinct responses
+    /// than the window holds can count a very old identity a second time
+    /// without recording it here. The window is far larger than the widest
+    /// repeat distance either provider produces. A repeat is usually dropped, but a Claude repeat
+    /// carrying larger counts than the copy already counted replaces it,
+    /// because a streaming response's early lines carry placeholder
+    /// output counts. Either way it counts here.
+    public var duplicateReports: Int
+
+    /// Usage blocks skipped because they carried no response identity.
+    ///
+    /// These are not counted, so a non-zero value means the total is an
+    /// undercount. That is the deliberate direction to be wrong in: without
+    /// an identity there is no way to tell a fresh response from the same
+    /// one reported again, and counting it risks the large overstatement
+    /// deduplication exists to prevent.
+    ///
+    /// Also counts a usage block that carries no count field this parser
+    /// recognizes, which is the one format change it can notice. In
+    /// practice both stay zero. It is not a general format alarm: a
+    /// provider that renames only *some* count keys, or adds a new kind of
+    /// token, leaves this at zero and quietly undercounts.
+    public var unidentifiedReports: Int
+
+    /// Whether cumulative-only Codex usage crossed an unmarked decrease.
+    ///
+    /// Without a structured thread/reset event, the accumulator cannot know
+    /// whether the smaller value corrected one run or began another. The
+    /// reported cumulative usage is then the latest known snapshot rather
+    /// than a guessed sum of runs.
+    public var cumulativeUsageIsAmbiguous: Bool
+
+    /// Tokens currently occupying the context window, when the provider
+    /// reports it.
+    public var contextTokens: Int?
+
+    /// The model's context window size, when the provider reports it.
+    public var contextWindowTokens: Int?
+
+    /// The provider's usage allowance state, when it reports it.
+    public var rateLimit: ChatUsageRateLimit?
+
+    /// Creates a usage summary.
+    ///
+    /// - Parameters:
+    ///   - usage: Usage over every distinct response.
+    ///   - usageByModel: Usage split by model identifier.
+    ///   - responses: Distinct responses counted.
+    ///   - duplicateReports: Repeated reports skipped.
+    ///   - unidentifiedReports: Usage blocks skipped for lack of an identity.
+    ///   - cumulativeUsageIsAmbiguous: Whether a cumulative-only transition
+    ///     lacked a structured run boundary.
+    ///   - contextTokens: Tokens currently in the context window.
+    ///   - contextWindowTokens: The context window size.
+    ///   - rateLimit: The provider's allowance state.
+    public init(
+        usage: ChatTokenUsage = ChatTokenUsage(),
+        usageByModel: [String: ChatTokenUsage] = [:],
+        responses: Int = 0,
+        duplicateReports: Int = 0,
+        unidentifiedReports: Int = 0,
+        cumulativeUsageIsAmbiguous: Bool = false,
+        contextTokens: Int? = nil,
+        contextWindowTokens: Int? = nil,
+        rateLimit: ChatUsageRateLimit? = nil
+    ) {
+        self.usage = usage
+        self.usageByModel = usageByModel
+        self.responses = responses
+        self.duplicateReports = duplicateReports
+        self.unidentifiedReports = unidentifiedReports
+        self.cumulativeUsageIsAmbiguous = cumulativeUsageIsAmbiguous
+        self.contextTokens = contextTokens
+        self.contextWindowTokens = contextWindowTokens
+        self.rateLimit = rateLimit
+    }
+
+    /// The share of the context window in use, 0 to 1, when both the
+    /// occupancy and the window size are known.
+    public var contextUsedFraction: Double? {
+        guard let contextTokens, let contextWindowTokens, contextWindowTokens > 0 else {
+            return nil
+        }
+        return Double(contextTokens) / Double(contextWindowTokens)
+    }
+}