@@ -34,11 +34,19 @@ class UsageViewModel: ObservableObject {
3434 case notConfigured
3535 }
3636
37+ let accountStore : AccountStore
3738 private let pollingInterval : TimeInterval = 300 // 5 minutes
3839 private var pollingTask : Task < Void , Never > ?
3940 private var cancellables = Set < AnyCancellable > ( )
4041
41- init ( ) {
42+ /// The email of the currently active account.
43+ var activeEmail : String ? {
44+ accountStore. activeAccount? . email
45+ }
46+
47+ init ( accountStore: AccountStore ) {
48+ self . accountStore = accountStore
49+
4250 // Initialize time display preference from UserDefaults
4351 let savedFormat = UserDefaults . standard. string ( forKey: " claude_time_display_format " ) ?? TimeDisplayFormat . resetTime. rawValue
4452 timeDisplayFormat = TimeDisplayFormat ( rawValue: savedFormat) ?? . resetTime
@@ -65,25 +73,52 @@ class UsageViewModel: ObservableObject {
6573 }
6674 . store ( in: & cancellables)
6775
68- // Start polling if credentials exist
69- startPollingIfConfigured ( )
76+ // Watch for active account changes and restart polling
77+ // Note: @Published fires on willSet, so we receive on next runloop tick
78+ // to ensure activeAccountId is already updated when we read it.
79+ accountStore. $activeAccountId
80+ . removeDuplicates ( )
81+ . receive ( on: RunLoop . main)
82+ . sink { [ weak self] _ in
83+ guard let self = self else { return }
84+ // Cancel existing polling FIRST to prevent stale fetches
85+ self . pollingTask? . cancel ( )
86+ self . pollingTask = nil
87+ self . usageData = nil
88+ self . lastUpdated = nil
89+ self . errorState = nil
90+ self . updateAuthStatus ( )
91+ self . startPollingIfConfigured ( )
92+ }
93+ . store ( in: & cancellables)
7094 }
7195
7296 deinit {
7397 pollingTask? . cancel ( )
7498 }
7599
76100 func updateAuthStatus( ) {
77- let hasKey = KeychainService . read ( key: . sessionKey) != nil
78- let hasOrg = KeychainService . read ( key: . orgId) != nil
101+ guard let account = accountStore. activeAccount else {
102+ authStatus = . notConfigured
103+ return
104+ }
105+ // Check both in-memory model and Keychain for credentials
106+ let hasKey : Bool = {
107+ if let key = account. sessionKey, !key. isEmpty { return true }
108+ if let key = accountStore. sessionKey ( for: account. id) , !key. isEmpty { return true }
109+ return false
110+ } ( )
111+ let hasOrg : Bool = {
112+ if let org = account. orgId, !org. isEmpty { return true }
113+ if let org = accountStore. orgId ( for: account. id) , !org. isEmpty { return true }
114+ return false
115+ } ( )
79116 if hasKey && hasOrg {
80117 if errorState == . authExpired {
81118 authStatus = . expired
82- } else if usageData != nil {
83- authStatus = . connected
84119 } else {
85- // Credentials exist but haven't verified yet
86- authStatus = . notConfigured
120+ // Credentials exist — treat as connected (data fetch in progress or complete)
121+ authStatus = . connected
87122 }
88123 } else {
89124 authStatus = . notConfigured
@@ -94,53 +129,58 @@ class UsageViewModel: ObservableObject {
94129 func startPollingIfConfigured( ) {
95130 pollingTask? . cancel ( )
96131
97- guard let sessionKey = KeychainService . read ( key: . sessionKey) ,
98- let orgId = KeychainService . read ( key: . orgId) ,
132+ guard let accountId = accountStore. activeAccountId,
133+ let sessionKey = accountStore. sessionKey ( for: accountId) ,
134+ let orgId = accountStore. orgId ( for: accountId) ,
99135 !sessionKey. isEmpty, !orgId. isEmpty else {
100136 return
101137 }
102138
103139 pollingTask = Task { [ weak self] in
104140 guard let self = self else { return }
105141 // Initial fetch
106- await self . performFetch ( sessionKey: sessionKey, orgId: orgId)
142+ guard !Task. isCancelled, self . accountStore. activeAccountId == accountId else { return }
143+ await self . performFetch ( sessionKey: sessionKey, orgId: orgId, accountId: accountId)
107144
108145 // Polling loop
109146 while !Task. isCancelled {
110147 try ? await Task . sleep ( nanoseconds: UInt64 ( self . pollingInterval * 1_000_000_000 ) )
111148 if Task . isCancelled { break }
149+ // Verify this is still the active account
150+ guard self . accountStore. activeAccountId == accountId else { break }
112151 // Re-read credentials in case they were updated
113- guard let currentKey = KeychainService . read ( key : . sessionKey ) ,
114- let currentOrg = KeychainService . read ( key : . orgId ) else {
152+ guard let currentKey = self . accountStore . sessionKey ( for : accountId ) ,
153+ let currentOrg = self . accountStore . orgId ( for : accountId ) else {
115154 break
116155 }
117- await self . performFetch ( sessionKey: currentKey, orgId: currentOrg)
156+ await self . performFetch ( sessionKey: currentKey, orgId: currentOrg, accountId : accountId )
118157 }
119158 }
120159 }
121160
122161 /// Triggers an immediate fetch outside the polling cycle.
123162 func fetchNow( ) {
124- guard let sessionKey = KeychainService . read ( key: . sessionKey) ,
125- let orgId = KeychainService . read ( key: . orgId) ,
163+ guard let accountId = accountStore. activeAccountId,
164+ let sessionKey = accountStore. sessionKey ( for: accountId) ,
165+ let orgId = accountStore. orgId ( for: accountId) ,
126166 !sessionKey. isEmpty, !orgId. isEmpty else {
127167 return
128168 }
129169
130170 Task { [ weak self] in
131- await self ? . performFetch ( sessionKey: sessionKey, orgId: orgId)
171+ await self ? . performFetch ( sessionKey: sessionKey, orgId: orgId, accountId : accountId )
132172 }
133173 }
134174
135175 @MainActor
136- private func performFetch( sessionKey: String , orgId: String ) async {
176+ private func performFetch( sessionKey: String , orgId: String , accountId : UUID ) async {
137177 let service = UsageService ( sessionKey: sessionKey, orgId: orgId)
138178 do {
139179 let ( data, newKey) = try await service. fetchUsage ( )
140180
141181 // Update Keychain if a new session key was returned via Set-Cookie
142182 if let newKey = newKey {
143- KeychainService . save ( key : . sessionKey , value : newKey )
183+ accountStore . saveSessionKey ( newKey , for : accountId )
144184 }
145185
146186 usageData = data
0 commit comments