From 8584ced7003ac779e87f9e8294e0ac66eae2cbb6 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Mon, 10 Aug 2026 12:24:41 -0700 Subject: [PATCH] perf: stop dormant agent session schedulers --- CHANGELOG.md | 1 + Sources/CodexBar/AgentSessionsStore.swift | 307 +++++++++++--- .../AgentSessionsStoreSchedulerTests.swift | 380 ++++++++++++++++++ 3 files changed, 624 insertions(+), 64 deletions(-) create mode 100644 Tests/CodexBarTests/AgentSessionsStoreSchedulerTests.swift diff --git a/CHANGELOG.md b/CHANGELOG.md index 671a07d1f8..4c8bc22515 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ - Sub2API: localize and group menu-card quota labels and request, token, and cost totals into a compact usage summary (#2835). Thanks @weirdo-adam! - Codex: avoid repeatedly converting historical token snapshots during cost-cache refreshes, preventing sustained CPU usage on large session histories. - Codex: make automatic cost-history catch-up near-idle and limit local-history scans to provider refreshes with a 15-minute energy floor. +- Agent Sessions: stop inactive local and remote refresh schedulers from waking while monitoring is disabled. ## 0.49.1 — 2026-08-09 diff --git a/Sources/CodexBar/AgentSessionsStore.swift b/Sources/CodexBar/AgentSessionsStore.swift index 45a9adec5a..7e817cefa9 100644 --- a/Sources/CodexBar/AgentSessionsStore.swift +++ b/Sources/CodexBar/AgentSessionsStore.swift @@ -2,7 +2,7 @@ import CodexBarCore import Foundation import Observation -struct AgentSessionRemoteRefreshGate { +struct AgentSessionRefreshGate { private(set) var generation = 0 private(set) var isInFlight = false private(set) var isPending = false @@ -29,20 +29,39 @@ struct AgentSessionRemoteRefreshGate { } } +typealias AgentSessionRemoteRefreshGate = AgentSessionRefreshGate + @MainActor @Observable final class AgentSessionsStore { typealias LocalScan = @Sendable (_ includeFileOnlySessions: Bool) async -> [AgentSession] + typealias RemoteHostDiscovery = @Sendable () async -> [String] + typealias RemoteFetch = @Sendable (_ hosts: [String]) async -> [RemoteSessionHostResult] + typealias PeriodicSleep = @Sendable (_ duration: Duration) async throws -> Void + + struct SchedulerState: Equatable { + let isStarted: Bool + let hasLocalPeriodicTask: Bool + let hasRemotePeriodicTask: Bool + let hasLocalImmediateTask: Bool + let hasRemoteImmediateTask: Bool + } private let settings: SettingsStore private let localScan: LocalScan + private let remoteHostDiscovery: RemoteHostDiscovery + private let remoteFetch: RemoteFetch private let remoteFetcher: RemoteSessionFetcher - @ObservationIgnored private var localRefreshTask: Task? - @ObservationIgnored private var remoteRefreshTask: Task? - @ObservationIgnored private var localRefreshInFlight = false + private let periodicSleep: PeriodicSleep + @ObservationIgnored private var localPeriodicTask: Task? + @ObservationIgnored private var remotePeriodicTask: Task? + @ObservationIgnored private var localImmediateTask: Task? + @ObservationIgnored private var remoteImmediateTask: Task? + @ObservationIgnored private var localRefreshGate = AgentSessionRefreshGate() @ObservationIgnored private var remoteRefreshGate = AgentSessionRemoteRefreshGate() @ObservationIgnored var onUpdate: (@MainActor () -> Void)? + private(set) var isStarted = false private(set) var localSessions: [AgentSession] = [] private(set) var remoteHosts: [RemoteSessionHostResult] = [] private(set) var lastUpdatedAt: Date? @@ -57,7 +76,16 @@ final class AgentSessionsStore { self.localScan = { includeFileOnlySessions in await localScanner.scan(includeFileOnlySessions: includeFileOnlySessions) } + self.remoteHostDiscovery = { + await remoteFetcher.discoveredHosts() + } + self.remoteFetch = { hosts in + await remoteFetcher.fetch(hosts: hosts) + } self.remoteFetcher = remoteFetcher + self.periodicSleep = { duration in + try await Task.sleep(for: duration) + } } init( @@ -67,7 +95,38 @@ final class AgentSessionsStore { { self.settings = settings self.localScan = localScan + self.remoteHostDiscovery = { + await remoteFetcher.discoveredHosts() + } + self.remoteFetch = { hosts in + await remoteFetcher.fetch(hosts: hosts) + } self.remoteFetcher = remoteFetcher + self.periodicSleep = { duration in + try await Task.sleep(for: duration) + } + } + + init( + settings: SettingsStore, + localScan: @escaping LocalScan, + remoteHostDiscovery: @escaping RemoteHostDiscovery, + remoteFetch: @escaping RemoteFetch, + periodicSleep: @escaping PeriodicSleep = { duration in try await Task.sleep(for: duration) }) + { + self.settings = settings + self.localScan = localScan + self.remoteHostDiscovery = remoteHostDiscovery + self.remoteFetch = remoteFetch + self.remoteFetcher = RemoteSessionFetcher() + self.periodicSleep = periodicSleep + } + + deinit { + self.localPeriodicTask?.cancel() + self.remotePeriodicTask?.cancel() + self.localImmediateTask?.cancel() + self.remoteImmediateTask?.cancel() } var totalCount: Int { @@ -80,6 +139,15 @@ final class AgentSessionsStore { self.settings.agentSessionsEnabled || self.settings.adaptiveActivityScanningEnabled } + var schedulerState: SchedulerState { + SchedulerState( + isStarted: self.isStarted, + hasLocalPeriodicTask: self.localPeriodicTask != nil, + hasRemotePeriodicTask: self.remotePeriodicTask != nil, + hasLocalImmediateTask: self.localImmediateTask != nil, + hasRemoteImmediateTask: self.remoteImmediateTask != nil) + } + nonisolated static func latestActivityAt(in sessions: [AgentSession]) -> Date? { sessions.compactMap(\.lastActivityAt).max() } @@ -98,60 +166,56 @@ final class AgentSessionsStore { } func start() { - guard self.localRefreshTask == nil, self.remoteRefreshTask == nil else { return } - self.localRefreshTask = Task { [weak self] in - while !Task.isCancelled { - await self?.refreshLocal() - try? await Task.sleep(for: .seconds(30)) - } - } - self.remoteRefreshTask = Task { [weak self] in - while !Task.isCancelled { - await self?.refreshRemote() - try? await Task.sleep(for: .seconds(60)) - } - } + guard !self.isStarted else { return } + self.isStarted = true + self.localRefreshGate.settingsDidChange() + self.remoteRefreshGate.settingsDidChange() + self.reconcilePeriodicTasks() + self.requestLocalRefresh() + self.requestRemoteRefresh() } func stop() { - self.localRefreshTask?.cancel() - self.remoteRefreshTask?.cancel() - self.localRefreshTask = nil - self.remoteRefreshTask = nil + guard self.isStarted || self.hasOwnedTasks else { return } + self.isStarted = false + self.localRefreshGate.settingsDidChange() + self.remoteRefreshGate.settingsDidChange() + self.cancelOwnedTasks() } func settingsDidChange(remoteConfigurationChanged: Bool = true) { + self.localRefreshGate.settingsDidChange() if remoteConfigurationChanged { self.remoteRefreshGate.settingsDidChange() } + + let hadVisibleSessions = !self.localSessions.isEmpty || !self.remoteHosts.isEmpty + let hadActivity = self.latestLocalActivityAt != nil if !self.settings.agentSessionsEnabled { // Adaptive keeps only the timestamp signal. Retained session paths and identities // remain scoped to the explicitly enabled Agent Sessions UI. self.localSessions = [] self.remoteHosts = [] } - guard self.localMonitoringEnabled else { + if !self.localMonitoringEnabled { self.latestLocalActivityAt = nil + } + + self.reconcilePeriodicTasks() + guard self.isStarted else { return } + if hadVisibleSessions || (hadActivity && !self.localMonitoringEnabled) { self.onUpdate?() - return } - guard !SettingsStore.isRunningTests else { return } - Task { [weak self] in - await self?.refreshLocal() - if remoteConfigurationChanged, self?.settings.agentSessionsEnabled == true { - await self?.refreshRemote() - } + self.requestLocalRefresh() + if remoteConfigurationChanged { + self.requestRemoteRefresh() } } func refreshOnMenuOpen() { - guard self.localMonitoringEnabled, !SettingsStore.isRunningTests else { return } - Task { [weak self] in - await self?.refreshLocal() - if self?.settings.agentSessionsEnabled == true { - await self?.refreshRemote() - } - } + guard self.isStarted else { return } + self.requestLocalRefresh() + self.requestRemoteRefresh() } func focus(_ session: AgentSession, remoteHost: String?) { @@ -165,19 +229,9 @@ final class AgentSessionsStore { } func refreshLocal() async { - guard self.localMonitoringEnabled, !self.localRefreshInFlight else { return } - let processInfo = ProcessInfo.processInfo - guard Self.shouldScanLocally( - agentSessionsEnabled: self.settings.agentSessionsEnabled, - adaptiveActivityScanningEnabled: self.settings.adaptiveActivityScanningEnabled, - lowPowerModeEnabled: processInfo.isLowPowerModeEnabled, - thermalState: processInfo.thermalState) - else { return } - self.localRefreshInFlight = true - let sessions = await self.localScan(self.settings.agentSessionsEnabled) - self.localRefreshInFlight = false - guard !Task.isCancelled, self.localMonitoringEnabled else { return } - self.applyLocalScanResult(sessions) + self.requestLocalRefresh() + let task = self.localImmediateTask + await task?.value } func applyLocalScanResult(_ sessions: [AgentSession], updatedAt: Date = Date()) { @@ -196,22 +250,147 @@ final class AgentSessionsStore { self.onUpdate?() } - private func refreshRemote() async { - guard self.settings.agentSessionsEnabled else { return } - guard var generation = self.remoteRefreshGate.begin() else { return } - while self.settings.agentSessionsEnabled { - var hosts = self.manualHosts - await hosts.append(contentsOf: self.remoteFetcher.discoveredHosts()) - let results = await self.remoteFetcher.fetch(hosts: hosts) - let outcome = self.remoteRefreshGate.finish(generation: generation) - guard !Task.isCancelled, self.settings.agentSessionsEnabled else { return } - if outcome.shouldPublish, results != self.remoteHosts { - self.remoteHosts = results - self.lastUpdatedAt = Date() - self.onUpdate?() + private var hasOwnedTasks: Bool { + self.localPeriodicTask != nil || self.remotePeriodicTask != nil || + self.localImmediateTask != nil || self.remoteImmediateTask != nil + } + + private func reconcilePeriodicTasks() { + let needsLocalScheduler = self.isStarted && self.localMonitoringEnabled + if needsLocalScheduler, self.localPeriodicTask == nil { + let periodicSleep = self.periodicSleep + self.localPeriodicTask = Task { [weak self] in + while !Task.isCancelled { + do { + try await periodicSleep(.seconds(30)) + } catch { + return + } + guard !Task.isCancelled else { return } + self?.requestLocalRefresh() + } + } + } else if !needsLocalScheduler { + self.localPeriodicTask?.cancel() + self.localPeriodicTask = nil + self.localImmediateTask?.cancel() + } + + let needsRemoteScheduler = self.isStarted && self.settings.agentSessionsEnabled + if needsRemoteScheduler, self.remotePeriodicTask == nil { + let periodicSleep = self.periodicSleep + self.remotePeriodicTask = Task { [weak self] in + while !Task.isCancelled { + do { + try await periodicSleep(.seconds(60)) + } catch { + return + } + guard !Task.isCancelled else { return } + self?.requestRemoteRefresh() + } } - guard outcome.shouldRetry, let nextGeneration = self.remoteRefreshGate.begin() else { return } - generation = nextGeneration + } else if !needsRemoteScheduler { + self.remotePeriodicTask?.cancel() + self.remotePeriodicTask = nil + self.remoteImmediateTask?.cancel() + } + } + + private func cancelOwnedTasks() { + self.localPeriodicTask?.cancel() + self.remotePeriodicTask?.cancel() + self.localImmediateTask?.cancel() + self.remoteImmediateTask?.cancel() + self.localPeriodicTask = nil + self.remotePeriodicTask = nil + self.localImmediateTask = nil + self.remoteImmediateTask = nil + } + + private func requestLocalRefresh() { + guard self.isStarted, self.localMonitoringEnabled, self.localImmediateTask == nil else { return } + let processInfo = ProcessInfo.processInfo + guard Self.shouldScanLocally( + agentSessionsEnabled: self.settings.agentSessionsEnabled, + adaptiveActivityScanningEnabled: self.settings.adaptiveActivityScanningEnabled, + lowPowerModeEnabled: processInfo.isLowPowerModeEnabled, + thermalState: processInfo.thermalState) + else { return } + guard let generation = self.localRefreshGate.begin() else { return } + + let includeFileOnlySessions = self.settings.agentSessionsEnabled + let localScan = self.localScan + self.localImmediateTask = Task { [weak self] in + guard !Task.isCancelled else { + self?.completeLocalRefresh(generation: generation, sessions: nil, wasCancelled: true) + return + } + let sessions = await localScan(includeFileOnlySessions) + self?.completeLocalRefresh( + generation: generation, + sessions: sessions, + wasCancelled: Task.isCancelled) + } + } + + private func completeLocalRefresh( + generation: Int, + sessions: [AgentSession]?, + wasCancelled: Bool) + { + self.localImmediateTask = nil + let outcome = self.localRefreshGate.finish(generation: generation) + if !wasCancelled, outcome.shouldPublish, self.isStarted, self.localMonitoringEnabled, let sessions { + self.applyLocalScanResult(sessions) + } + if outcome.shouldRetry, self.isStarted, self.localMonitoringEnabled { + self.requestLocalRefresh() + } + } + + private func requestRemoteRefresh() { + guard self.isStarted, self.settings.agentSessionsEnabled, self.remoteImmediateTask == nil else { return } + guard let generation = self.remoteRefreshGate.begin() else { return } + + let manualHosts = self.manualHosts + let remoteHostDiscovery = self.remoteHostDiscovery + let remoteFetch = self.remoteFetch + self.remoteImmediateTask = Task { [weak self] in + guard !Task.isCancelled else { + self?.completeRemoteRefresh(generation: generation, results: nil, wasCancelled: true) + return + } + var hosts = manualHosts + await hosts.append(contentsOf: remoteHostDiscovery()) + let results = await remoteFetch(hosts) + self?.completeRemoteRefresh( + generation: generation, + results: results, + wasCancelled: Task.isCancelled) + } + } + + private func completeRemoteRefresh( + generation: Int, + results: [RemoteSessionHostResult]?, + wasCancelled: Bool) + { + self.remoteImmediateTask = nil + let outcome = self.remoteRefreshGate.finish(generation: generation) + if !wasCancelled, + outcome.shouldPublish, + self.isStarted, + self.settings.agentSessionsEnabled, + let results, + results != self.remoteHosts + { + self.remoteHosts = results + self.lastUpdatedAt = Date() + self.onUpdate?() + } + if outcome.shouldRetry, self.isStarted, self.settings.agentSessionsEnabled { + self.requestRemoteRefresh() } } diff --git a/Tests/CodexBarTests/AgentSessionsStoreSchedulerTests.swift b/Tests/CodexBarTests/AgentSessionsStoreSchedulerTests.swift new file mode 100644 index 0000000000..fc7c8d8deb --- /dev/null +++ b/Tests/CodexBarTests/AgentSessionsStoreSchedulerTests.swift @@ -0,0 +1,380 @@ +import CodexBarCore +import Foundation +import Testing +@testable import CodexBar + +private actor AgentSessionScanHarness { + struct Call: Equatable, Sendable { + let includeFileOnlySessions: Bool + } + + private var calls: [Call] = [] + private var continuations: [Int: CheckedContinuation<[AgentSession], Never>] = [:] + private var callWaiters: [(count: Int, continuation: CheckedContinuation)] = [] + + func scan(includeFileOnlySessions: Bool) async -> [AgentSession] { + let index = self.calls.count + self.calls.append(Call(includeFileOnlySessions: includeFileOnlySessions)) + self.resumeSatisfiedWaiters() + return await withCheckedContinuation { continuation in + self.continuations[index] = continuation + } + } + + func waitForCallCount(_ count: Int) async { + guard self.calls.count < count else { return } + await withCheckedContinuation { continuation in + self.callWaiters.append((count, continuation)) + } + } + + func releaseCall(_ index: Int, returning sessions: [AgentSession]) { + self.continuations.removeValue(forKey: index)?.resume(returning: sessions) + } + + func recordedCalls() -> [Call] { + self.calls + } + + private func resumeSatisfiedWaiters() { + var remaining: [(count: Int, continuation: CheckedContinuation)] = [] + for waiter in self.callWaiters { + if self.calls.count >= waiter.count { + waiter.continuation.resume() + } else { + remaining.append(waiter) + } + } + self.callWaiters = remaining + } +} + +private actor AgentSessionRemoteFetchHarness { + private var calls: [[String]] = [] + private var continuations: [Int: CheckedContinuation<[RemoteSessionHostResult], Never>] = [:] + private var callWaiters: [(count: Int, continuation: CheckedContinuation)] = [] + + func fetch(hosts: [String]) async -> [RemoteSessionHostResult] { + let index = self.calls.count + self.calls.append(hosts) + self.resumeSatisfiedWaiters() + return await withCheckedContinuation { continuation in + self.continuations[index] = continuation + } + } + + func waitForCallCount(_ count: Int) async { + guard self.calls.count < count else { return } + await withCheckedContinuation { continuation in + self.callWaiters.append((count, continuation)) + } + } + + func releaseCall(_ index: Int, returning results: [RemoteSessionHostResult]) { + self.continuations.removeValue(forKey: index)?.resume(returning: results) + } + + func recordedCalls() -> [[String]] { + self.calls + } + + private func resumeSatisfiedWaiters() { + var remaining: [(count: Int, continuation: CheckedContinuation)] = [] + for waiter in self.callWaiters { + if self.calls.count >= waiter.count { + waiter.continuation.resume() + } else { + remaining.append(waiter) + } + } + self.callWaiters = remaining + } +} + +private actor AgentSessionRefreshSpy { + private(set) var localCalls: [Bool] = [] + private(set) var remoteCalls: [[String]] = [] + private var localWaiters: [(count: Int, continuation: CheckedContinuation)] = [] + private var remoteWaiters: [(count: Int, continuation: CheckedContinuation)] = [] + + func scan(includeFileOnlySessions: Bool) -> [AgentSession] { + self.localCalls.append(includeFileOnlySessions) + self.resumeLocalWaiters() + return [] + } + + func fetch(hosts: [String]) -> [RemoteSessionHostResult] { + self.remoteCalls.append(hosts) + self.resumeRemoteWaiters() + return [] + } + + func waitForLocalCallCount(_ count: Int) async { + guard self.localCalls.count < count else { return } + await withCheckedContinuation { continuation in + self.localWaiters.append((count, continuation)) + } + } + + func waitForRemoteCallCount(_ count: Int) async { + guard self.remoteCalls.count < count else { return } + await withCheckedContinuation { continuation in + self.remoteWaiters.append((count, continuation)) + } + } + + private func resumeLocalWaiters() { + let ready = self.localWaiters.filter { self.localCalls.count >= $0.count } + self.localWaiters.removeAll { self.localCalls.count >= $0.count } + ready.forEach { $0.continuation.resume() } + } + + private func resumeRemoteWaiters() { + let ready = self.remoteWaiters.filter { self.remoteCalls.count >= $0.count } + self.remoteWaiters.removeAll { self.remoteCalls.count >= $0.count } + ready.forEach { $0.continuation.resume() } + } +} + +@MainActor +struct AgentSessionsStoreSchedulerTests { + @Test + func `periodic schedulers match all lifecycle and settings states`() { + let settings = testSettingsStore(suiteName: "AgentSessionsStoreSchedulerTests-states") + let store = Self.makeStore(settings: settings) + + #expect(store.schedulerState == .init( + isStarted: false, + hasLocalPeriodicTask: false, + hasRemotePeriodicTask: false, + hasLocalImmediateTask: false, + hasRemoteImmediateTask: false)) + + store.start() + #expect(store.isStarted) + #expect(!store.schedulerState.hasLocalPeriodicTask) + #expect(!store.schedulerState.hasRemotePeriodicTask) + + settings.refreshFrequency = .adaptiveAgentAware + settings.adaptiveActivityScanConsent = .allowed + store.settingsDidChange(remoteConfigurationChanged: false) + #expect(store.schedulerState.hasLocalPeriodicTask) + #expect(!store.schedulerState.hasRemotePeriodicTask) + + settings.agentSessionsEnabled = true + store.settingsDidChange() + #expect(store.schedulerState.hasLocalPeriodicTask) + #expect(store.schedulerState.hasRemotePeriodicTask) + + store.stop() + #expect(store.schedulerState == .init( + isStarted: false, + hasLocalPeriodicTask: false, + hasRemotePeriodicTask: false, + hasLocalImmediateTask: false, + hasRemoteImmediateTask: false)) + } + + @Test + func `settings changes immediately enable and disable the owned schedulers`() async { + let settings = testSettingsStore(suiteName: "AgentSessionsStoreSchedulerTests-enable-disable") + let spy = AgentSessionRefreshSpy() + let store = Self.makeStore(settings: settings, spy: spy) + store.start() + + settings.refreshFrequency = .adaptiveAgentAware + settings.adaptiveActivityScanConsent = .allowed + store.settingsDidChange(remoteConfigurationChanged: false) + await spy.waitForLocalCallCount(1) + #expect(store.schedulerState.hasLocalPeriodicTask) + #expect(!store.schedulerState.hasRemotePeriodicTask) + #expect(await spy.localCalls == [false]) + + settings.agentSessionsEnabled = true + store.settingsDidChange() + await spy.waitForLocalCallCount(2) + await spy.waitForRemoteCallCount(1) + #expect(store.schedulerState.hasLocalPeriodicTask) + #expect(store.schedulerState.hasRemotePeriodicTask) + #expect(await spy.localCalls == [false, true]) + + settings.agentSessionsEnabled = false + settings.adaptiveActivityScanConsent = .declined + store.settingsDidChange() + #expect(!store.schedulerState.hasLocalPeriodicTask) + #expect(!store.schedulerState.hasRemotePeriodicTask) + #expect(store.localSessions.isEmpty) + #expect(store.remoteHosts.isEmpty) + #expect(store.latestLocalActivityAt == nil) + store.stop() + } + + @Test + func `rapid off then agent sessions transition rejects stale local data and retries once`() async { + let settings = testSettingsStore(suiteName: "AgentSessionsStoreSchedulerTests-local-generation") + settings.refreshFrequency = .adaptiveAgentAware + settings.adaptiveActivityScanConsent = .allowed + let scan = AgentSessionScanHarness() + let store = Self.makeStore(settings: settings, scan: scan) + store.start() + await scan.waitForCallCount(1) + + settings.adaptiveActivityScanConsent = .declined + store.settingsDidChange(remoteConfigurationChanged: false) + #expect(!store.schedulerState.hasLocalPeriodicTask) + + settings.agentSessionsEnabled = true + store.settingsDidChange() + await scan.releaseCall(0, returning: [Self.session(id: "stale", activity: Date(timeIntervalSince1970: 1))]) + await scan.waitForCallCount(2) + + #expect(await scan.recordedCalls() == [ + .init(includeFileOnlySessions: false), + .init(includeFileOnlySessions: true), + ]) + #expect(store.localSessions.isEmpty) + #expect(store.latestLocalActivityAt == nil) + + let current = Self.session(id: "current", activity: Date(timeIntervalSince1970: 2)) + await scan.releaseCall(1, returning: [current]) + await Self.waitForImmediateTasksToFinish(store) + #expect(store.localSessions == [current]) + #expect(store.latestLocalActivityAt == current.lastActivityAt) + #expect(await scan.recordedCalls().count == 2) + store.stop() + } + + @Test + func `remote configuration change rejects stale data and retries with current hosts`() async { + let settings = testSettingsStore(suiteName: "AgentSessionsStoreSchedulerTests-remote-generation") + settings.agentSessionsEnabled = true + settings.agentSessionsManualHosts = "old-host" + let remote = AgentSessionRemoteFetchHarness() + let store = Self.makeStore(settings: settings, remote: remote) + store.start() + await remote.waitForCallCount(1) + + settings.agentSessionsManualHosts = "new-host" + store.settingsDidChange() + await remote.releaseCall(0, returning: [ + RemoteSessionHostResult(host: "old-host", sessions: [], error: nil), + ]) + await remote.waitForCallCount(2) + + #expect(await remote.recordedCalls() == [["old-host"], ["new-host"]]) + #expect(store.remoteHosts.isEmpty) + + let current = RemoteSessionHostResult(host: "new-host", sessions: [], error: nil) + await remote.releaseCall(1, returning: [current]) + await Self.waitForImmediateTasksToFinish(store) + #expect(store.remoteHosts == [current]) + store.stop() + } + + @Test + func `menu refresh overlaps coalesce without an extra local pass`() async { + let settings = testSettingsStore(suiteName: "AgentSessionsStoreSchedulerTests-menu-coalescing") + settings.agentSessionsEnabled = true + let scan = AgentSessionScanHarness() + let store = Self.makeStore(settings: settings, scan: scan) + store.start() + await scan.waitForCallCount(1) + + for _ in 0..<5 { + store.refreshOnMenuOpen() + } + #expect(await scan.recordedCalls().count == 1) + + await scan.releaseCall(0, returning: []) + await Self.waitForImmediateTasksToFinish(store) + #expect(await scan.recordedCalls().count == 1) + store.stop() + } + + @Test + func `stop cancels every owned task and blocks late publication`() async { + let settings = testSettingsStore(suiteName: "AgentSessionsStoreSchedulerTests-stop") + settings.refreshFrequency = .adaptiveAgentAware + settings.adaptiveActivityScanConsent = .allowed + let scan = AgentSessionScanHarness() + let store = Self.makeStore(settings: settings, scan: scan) + var updateCount = 0 + store.onUpdate = { updateCount += 1 } + store.start() + await scan.waitForCallCount(1) + + store.stop() + #expect(!store.isStarted) + #expect(!store.schedulerState.hasLocalPeriodicTask) + #expect(!store.schedulerState.hasRemotePeriodicTask) + #expect(!store.schedulerState.hasLocalImmediateTask) + #expect(!store.schedulerState.hasRemoteImmediateTask) + + await scan.releaseCall(0, returning: [Self.session(id: "late", activity: Date())]) + for _ in 0..<20 { + await Task.yield() + } + #expect(store.localSessions.isEmpty) + #expect(store.latestLocalActivityAt == nil) + #expect(updateCount == 0) + } + + private static func makeStore( + settings: SettingsStore, + scan: AgentSessionScanHarness? = nil, + remote: AgentSessionRemoteFetchHarness? = nil) -> AgentSessionsStore + { + AgentSessionsStore( + settings: settings, + localScan: { includeFileOnlySessions in + guard let scan else { return [] } + return await scan.scan(includeFileOnlySessions: includeFileOnlySessions) + }, + remoteHostDiscovery: { [] }, + remoteFetch: { hosts in + guard let remote else { return [] } + return await remote.fetch(hosts: hosts) + }) + } + + private static func makeStore( + settings: SettingsStore, + spy: AgentSessionRefreshSpy) -> AgentSessionsStore + { + AgentSessionsStore( + settings: settings, + localScan: { includeFileOnlySessions in + await spy.scan(includeFileOnlySessions: includeFileOnlySessions) + }, + remoteHostDiscovery: { [] }, + remoteFetch: { hosts in + await spy.fetch(hosts: hosts) + }) + } + + private static func waitForImmediateTasksToFinish(_ store: AgentSessionsStore) async { + for _ in 0..<1000 { + let state = store.schedulerState + if !state.hasLocalImmediateTask, !state.hasRemoteImmediateTask { + return + } + await Task.yield() + } + Issue.record("Agent session immediate refresh tasks did not finish") + } + + private static func session(id: String, activity: Date?) -> AgentSession { + AgentSession( + id: id, + provider: .codex, + source: .cli, + state: .active, + pid: 42, + cwd: "/Users/test/alpha", + projectName: "alpha", + startedAt: nil, + lastActivityAt: activity, + transcriptPath: nil, + host: "local") + } +}