diff --git a/OrbitDockNative/OrbitDock/Models/RootShell/RootSessionNode.swift b/OrbitDockNative/OrbitDock/Models/RootShell/RootSessionNode.swift index 9e4b7078..cf63a4c4 100644 --- a/OrbitDockNative/OrbitDock/Models/RootShell/RootSessionNode.swift +++ b/OrbitDockNative/OrbitDock/Models/RootShell/RootSessionNode.swift @@ -148,6 +148,9 @@ struct RootSessionNode: Identifiable, Sendable { let missionId: String? let issueIdentifier: String? let totalTokens: Int + let inputTokens: Int + let outputTokens: Int + let cachedTokens: Int let totalCostUSD: Double let isActive: Bool let showsInMissionControl: Bool @@ -281,6 +284,9 @@ extension RootSessionNode: Equatable { && lhs.missionId == rhs.missionId && lhs.issueIdentifier == rhs.issueIdentifier && lhs.totalTokens == rhs.totalTokens + && lhs.inputTokens == rhs.inputTokens + && lhs.outputTokens == rhs.outputTokens + && lhs.cachedTokens == rhs.cachedTokens && lhs.totalCostUSD == rhs.totalCostUSD && lhs.isActive == rhs.isActive && lhs.showsInMissionControl == rhs.showsInMissionControl @@ -376,6 +382,9 @@ extension RootSessionNode { self.missionId = session.missionId self.issueIdentifier = session.issueIdentifier self.totalTokens = Int(session.totalTokens ?? 0) + self.inputTokens = Int(session.inputTokens ?? 0) + self.outputTokens = Int(session.outputTokens ?? 0) + self.cachedTokens = Int(session.cachedTokens ?? 0) self.totalCostUSD = session.totalCostUSD ?? 0 self.isActive = status == .active self.showsInMissionControl = RootSessionNode.showsInMissionControl( @@ -435,6 +444,9 @@ extension RootSessionNode { missionId: missionId, issueIdentifier: issueIdentifier, totalTokens: totalTokens, + inputTokens: inputTokens, + outputTokens: outputTokens, + cachedTokens: cachedTokens, totalCostUSD: totalCostUSD, isActive: self.status == .active, showsInMissionControl: RootSessionNode.showsInMissionControl( @@ -486,6 +498,9 @@ extension RootSessionNode { missionId: missionId, issueIdentifier: issueIdentifier, totalTokens: totalTokens, + inputTokens: inputTokens, + outputTokens: outputTokens, + cachedTokens: cachedTokens, totalCostUSD: totalCostUSD, isActive: false, showsInMissionControl: RootSessionNode.showsInMissionControl( diff --git a/OrbitDockNative/OrbitDock/Services/Server/Protocol/ServerSessionContracts.swift b/OrbitDockNative/OrbitDock/Services/Server/Protocol/ServerSessionContracts.swift index c5b11fa6..1a46dca7 100644 --- a/OrbitDockNative/OrbitDock/Services/Server/Protocol/ServerSessionContracts.swift +++ b/OrbitDockNative/OrbitDock/Services/Server/Protocol/ServerSessionContracts.swift @@ -207,6 +207,9 @@ struct ServerSessionListItem: Codable, Identifiable { let worktreeId: String? let totalTokens: UInt64? let totalCostUSD: Double? + let inputTokens: UInt64? + let outputTokens: UInt64? + let cachedTokens: UInt64? let displayTitle: String? let displayTitleSortKey: String? let displaySearchText: String? @@ -241,6 +244,9 @@ struct ServerSessionListItem: Codable, Identifiable { worktreeId: String?, totalTokens: UInt64?, totalCostUSD: Double?, + inputTokens: UInt64? = nil, + outputTokens: UInt64? = nil, + cachedTokens: UInt64? = nil, displayTitle: String?, displayTitleSortKey: String?, displaySearchText: String?, @@ -274,6 +280,9 @@ struct ServerSessionListItem: Codable, Identifiable { self.worktreeId = worktreeId self.totalTokens = totalTokens self.totalCostUSD = totalCostUSD + self.inputTokens = inputTokens + self.outputTokens = outputTokens + self.cachedTokens = cachedTokens self.displayTitle = displayTitle self.displayTitleSortKey = displayTitleSortKey self.displaySearchText = displaySearchText @@ -309,6 +318,9 @@ struct ServerSessionListItem: Codable, Identifiable { case worktreeId = "worktree_id" case totalTokens = "total_tokens" case totalCostUSD = "total_cost_usd" + case inputTokens = "input_tokens" + case outputTokens = "output_tokens" + case cachedTokens = "cached_tokens" case displayTitle = "display_title" case displayTitleSortKey = "display_title_sort_key" case displaySearchText = "display_search_text" diff --git a/OrbitDockNative/OrbitDock/Views/Dashboard/Usage/StatsPopoverContent.swift b/OrbitDockNative/OrbitDock/Views/Dashboard/Usage/StatsPopoverContent.swift index 0beb11ea..12e33378 100644 --- a/OrbitDockNative/OrbitDock/Views/Dashboard/Usage/StatsPopoverContent.swift +++ b/OrbitDockNative/OrbitDock/Views/Dashboard/Usage/StatsPopoverContent.swift @@ -273,6 +273,18 @@ struct StatusBarStats { return session.totalCostUSD } guard session.totalTokens > 0 else { return 0 } + // Use granular token breakdown when available for accurate cost. + // Cached tokens are billed at the cache-read rate, not the full input rate. + let hasBreakdown = session.inputTokens > 0 || session.outputTokens > 0 + if hasBreakdown { + return costCalculator.calculateCost( + model: session.model, + inputTokens: session.inputTokens, + outputTokens: session.outputTokens, + cacheReadTokens: session.cachedTokens + ) + } + // Legacy fallback: treat totalTokens as input (pre-breakdown servers). return costCalculator.calculateCost( model: session.model, inputTokens: session.totalTokens, diff --git a/OrbitDockNative/OrbitDock/Views/QuickSwitcher/QuickSwitcherPreview.swift b/OrbitDockNative/OrbitDock/Views/QuickSwitcher/QuickSwitcherPreview.swift index 58d68730..010adee6 100644 --- a/OrbitDockNative/OrbitDock/Views/QuickSwitcher/QuickSwitcherPreview.swift +++ b/OrbitDockNative/OrbitDock/Views/QuickSwitcher/QuickSwitcherPreview.swift @@ -86,6 +86,9 @@ private func quickSwitcherPreviewNode( missionId: nil, issueIdentifier: nil, totalTokens: 0, + inputTokens: 0, + outputTokens: 0, + cachedTokens: 0, totalCostUSD: 0, isActive: status == .active, showsInMissionControl: status == .active, diff --git a/OrbitDockNative/OrbitDockTests/SessionDetail/StatusBarStatsTests.swift b/OrbitDockNative/OrbitDockTests/SessionDetail/StatusBarStatsTests.swift index 4867b10f..484218a2 100644 --- a/OrbitDockNative/OrbitDockTests/SessionDetail/StatusBarStatsTests.swift +++ b/OrbitDockNative/OrbitDockTests/SessionDetail/StatusBarStatsTests.swift @@ -75,10 +75,73 @@ struct StatusBarStatsTests { #expect(stats.costByModel.first?.cost == stats.cost) } + @Test func fromUsesGranularTokenBreakdownForCost() { + // When input/output breakdown is available, cost should reflect + // the different rates for input vs output tokens. + let inputCount = 500_000 + let outputCount = 500_000 + let cachedCount = 100_000 + + let sessions = [ + makeSession( + id: "claude-1", + model: "claude-sonnet-4", + totalTokens: inputCount + outputCount, + inputTokens: inputCount, + outputTokens: outputCount, + cachedTokens: cachedCount, + totalCostUSD: 0 + ), + ] + + let stats = StatusBarStats.from( + sessions: sessions, + costCalculator: costCalculator + ) + + // With the fallback calculator (no prices loaded), defaults are: + // input: $3/M, output: $15/M, cache_read: $0.30/M + let expectedInput = Double(inputCount) / 1_000_000 * 3.0 + let expectedOutput = Double(outputCount) / 1_000_000 * 15.0 + let expectedCache = Double(cachedCount) / 1_000_000 * 0.30 + let expectedCost = expectedInput + expectedOutput + expectedCache + + #expect(stats.tokens == inputCount + outputCount) + #expect(abs(stats.cost - expectedCost) < 0.001) + } + + @Test func fromFallsBackToLegacyWhenNoBreakdown() { + // When input/output are both 0 (legacy server), falls back to + // treating totalTokens as input. + let sessions = [ + makeSession( + id: "claude-1", + model: "claude-sonnet-4", + totalTokens: 1_000_000, + inputTokens: 0, + outputTokens: 0, + cachedTokens: 0, + totalCostUSD: 0 + ), + ] + + let stats = StatusBarStats.from( + sessions: sessions, + costCalculator: costCalculator + ) + + // Legacy fallback: all tokens treated as input at $3/M + let expectedCost = Double(1_000_000) / 1_000_000 * 3.0 + #expect(abs(stats.cost - expectedCost) < 0.001) + } + private func makeSession( id: String, model: String, totalTokens: Int, + inputTokens: Int = 0, + outputTokens: Int = 0, + cachedTokens: Int = 0, totalCostUSD: Double ) -> RootSessionNode { RootSessionNode( @@ -116,6 +179,9 @@ struct StatusBarStatsTests { missionId: nil, issueIdentifier: nil, totalTokens: totalTokens, + inputTokens: inputTokens, + outputTokens: outputTokens, + cachedTokens: cachedTokens, totalCostUSD: totalCostUSD, isActive: true, showsInMissionControl: true, diff --git a/orbitdock-server/crates/protocol/src/types.rs b/orbitdock-server/crates/protocol/src/types.rs index 67038e12..104d0913 100644 --- a/orbitdock-server/crates/protocol/src/types.rs +++ b/orbitdock-server/crates/protocol/src/types.rs @@ -776,6 +776,9 @@ impl SessionSummary { worktree_id: self.worktree_id.clone(), total_tokens: self.token_usage.input_tokens + self.token_usage.output_tokens, total_cost_usd: 0.0, + input_tokens: self.token_usage.input_tokens, + output_tokens: self.token_usage.output_tokens, + cached_tokens: self.token_usage.cached_tokens, display_title: self.display_title.clone(), context_line: self.context_line.clone(), list_status: self.list_status, @@ -813,6 +816,9 @@ impl From for SessionListItem { worktree_id: summary.worktree_id, total_tokens: summary.token_usage.input_tokens + summary.token_usage.output_tokens, total_cost_usd: 0.0, + input_tokens: summary.token_usage.input_tokens, + output_tokens: summary.token_usage.output_tokens, + cached_tokens: summary.token_usage.cached_tokens, display_title: summary.display_title, context_line: summary.context_line, list_status: summary.list_status, @@ -1056,6 +1062,13 @@ pub struct SessionListItem { pub total_tokens: u64, #[serde(default)] pub total_cost_usd: f64, + /// Granular token breakdown for accurate cost calculation. + #[serde(default)] + pub input_tokens: u64, + #[serde(default)] + pub output_tokens: u64, + #[serde(default)] + pub cached_tokens: u64, pub display_title: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub context_line: Option, @@ -1105,6 +1118,9 @@ impl SessionListItem { worktree_id: summary.worktree_id.clone(), total_tokens: summary.token_usage.input_tokens + summary.token_usage.output_tokens, total_cost_usd: 0.0, + input_tokens: summary.token_usage.input_tokens, + output_tokens: summary.token_usage.output_tokens, + cached_tokens: summary.token_usage.cached_tokens, display_title: summary.display_title.clone(), context_line: summary.context_line.clone(), list_status: summary.list_status,