From 53cd7d2d12f6b8047cd9344c372c5c261911bc3b Mon Sep 17 00:00:00 2001 From: Dima Kopfmann Date: Mon, 29 Jun 2026 13:59:03 +0200 Subject: [PATCH] IJPL-246329 refresh working thread costs periodically Active agent threads can keep producing paid usage without changing the sessions tree state, so visible cost hydration must not depend solely on state-store emissions after the initial load. Keep the existing cache TTL semantics for finished threads, but add a periodic visible hydration wake-up while a visible thread is still working. Tests inject shorter intervals and cover a working thread whose cost changes without an explicit refresh. (cherry picked from commit 4cd9929a5993c4ee2a9428832920cb0d41df81ec) IJ-MR-210801 GitOrigin-RevId: 591bf79d529e88558c089df5f5e276e1ec6a2728 --- .../src/service/AgentSessionRefreshService.kt | 4 ++ ...AgentSessionVisibleCostHydrationSupport.kt | 34 ++++++++-- .../AgentSessionIntegrationTestSupport.kt | 20 ++++++ ...entSessionRefreshServiceIntegrationTest.kt | 63 +++++++++++++++++++ 4 files changed, 116 insertions(+), 5 deletions(-) diff --git a/plugins/agent-workbench/sessions/src/service/AgentSessionRefreshService.kt b/plugins/agent-workbench/sessions/src/service/AgentSessionRefreshService.kt index f1f63e3dbbc0..53b22ba4aa37 100644 --- a/plugins/agent-workbench/sessions/src/service/AgentSessionRefreshService.kt +++ b/plugins/agent-workbench/sessions/src/service/AgentSessionRefreshService.kt @@ -83,6 +83,8 @@ class AgentSessionRefreshService internal constructor( private val currentTimeMillis: () -> Long = System::currentTimeMillis, private val archiveTransitionSuppressions: AgentSessionArchiveTransitionSuppressions = AgentSessionArchiveTransitionSuppressions(), private val loadingDelayMs: Long = DEFAULT_AGENT_SESSION_LOADING_DELAY_MS, + private val visibleCostHydrationDelayMs: Long = 750L, + private val workingThreadCostCacheTtlMs: Long = 60_000L, ) { @Suppress("unused") constructor(serviceScope: CoroutineScope) : this( @@ -130,6 +132,8 @@ class AgentSessionRefreshService internal constructor( sessionSourcesProvider = sessionSourcesProvider, toolWindowVisibleFlow = toolWindowVisibleFlow, currentTimeMillis = currentTimeMillis, + visibleCostHydrationDelayMs = visibleCostHydrationDelayMs, + workingThreadCostCacheTtlMs = workingThreadCostCacheTtlMs, ) init { diff --git a/plugins/agent-workbench/sessions/src/service/AgentSessionVisibleCostHydrationSupport.kt b/plugins/agent-workbench/sessions/src/service/AgentSessionVisibleCostHydrationSupport.kt index 88e1a5648635..835a50a47887 100644 --- a/plugins/agent-workbench/sessions/src/service/AgentSessionVisibleCostHydrationSupport.kt +++ b/plugins/agent-workbench/sessions/src/service/AgentSessionVisibleCostHydrationSupport.kt @@ -26,8 +26,8 @@ import java.util.concurrent.ConcurrentHashMap import kotlin.time.Duration.Companion.milliseconds private val LOG = logger() -private val VISIBLE_COST_HYDRATION_DELAY = 750.milliseconds -private const val WORKING_THREAD_COST_CACHE_TTL_MS = 60_000L +private const val DEFAULT_VISIBLE_COST_HYDRATION_DELAY_MS = 750L +private const val DEFAULT_WORKING_THREAD_COST_CACHE_TTL_MS = 60_000L internal class AgentSessionVisibleCostHydrationSupport( private val serviceScope: CoroutineScope, @@ -36,6 +36,8 @@ internal class AgentSessionVisibleCostHydrationSupport( private val sessionSourcesProvider: () -> List, private val toolWindowVisibleFlow: StateFlow, private val currentTimeMillis: () -> Long = System::currentTimeMillis, + private val visibleCostHydrationDelayMs: Long = DEFAULT_VISIBLE_COST_HYDRATION_DELAY_MS, + private val workingThreadCostCacheTtlMs: Long = DEFAULT_WORKING_THREAD_COST_CACHE_TTL_MS, ) { private val costCache = ConcurrentHashMap() private val inFlightLoads = ConcurrentHashMap.newKeySet() @@ -52,12 +54,33 @@ internal class AgentSessionVisibleCostHydrationSupport( if (!snapshot.toolWindowVisible || !snapshot.costEnabled) { return@collectLatest } - delay(VISIBLE_COST_HYDRATION_DELAY) + delay(visibleCostHydrationDelayMs.milliseconds) hydrateVisibleThreadCosts(snapshot.state) + + while (hasVisibleWorkingThreadCostsToHydrate(stateStore.state.value)) { + delay(workingThreadCostCacheTtlMs.milliseconds) + if (!toolWindowVisibleFlow.value || !AgentSessionCostPresentationSettings.isEnabled()) { + return@collectLatest + } + hydrateVisibleThreadCosts(stateStore.state.value) + } } } } + private fun hasVisibleWorkingThreadCostsToHydrate(state: AgentSessionsState): Boolean { + val costSourceProviders = sessionSourcesProvider() + .asSequence() + .mapNotNull { source -> source as? AgentSessionCostSource } + .map { source -> source.provider } + .toSet() + return costSourceProviders.isNotEmpty() && collectVisibleThreads(state).any { visibleThread -> + visibleThread.provider in costSourceProviders && + visibleThread.activity.isWorking && + normalizeConcreteAgentSessionThreadId(visibleThread.threadId) != null + } + } + private fun hydrateVisibleThreadCosts(state: AgentSessionsState) { val visibleThreads = collectVisibleThreads(state) if (visibleThreads.isEmpty()) { @@ -89,7 +112,7 @@ internal class AgentSessionVisibleCostHydrationSupport( } val effectiveCacheEntry = costCache[cacheKey] - if (effectiveCacheEntry != null && shouldReuseCachedCost(visibleThread, effectiveCacheEntry, now)) { + if (effectiveCacheEntry != null && shouldReuseCachedCost(visibleThread, effectiveCacheEntry, now, workingThreadCostCacheTtlMs)) { if (visibleThread.cost != effectiveCacheEntry.cost) { cachedUpdatesByPath .getOrPut(visibleThread.path) { LinkedHashMap() } @@ -212,6 +235,7 @@ private fun shouldReuseCachedCost( visibleThread: VisibleThreadSnapshot, cacheEntry: ThreadCostCacheEntry, nowMs: Long, + workingThreadCostCacheTtlMs: Long, ): Boolean { if (cacheEntry.updatedAt != visibleThread.updatedAt) { return false @@ -219,7 +243,7 @@ private fun shouldReuseCachedCost( if (!visibleThread.activity.isWorking) { return true } - return nowMs - cacheEntry.refreshedAtMs < WORKING_THREAD_COST_CACHE_TTL_MS + return nowMs - cacheEntry.refreshedAtMs < workingThreadCostCacheTtlMs } internal data class ThreadCostUpdate( diff --git a/plugins/agent-workbench/sessions/testSrc/AgentSessionIntegrationTestSupport.kt b/plugins/agent-workbench/sessions/testSrc/AgentSessionIntegrationTestSupport.kt index 4b73a57cdc01..97419fcb99bb 100644 --- a/plugins/agent-workbench/sessions/testSrc/AgentSessionIntegrationTestSupport.kt +++ b/plugins/agent-workbench/sessions/testSrc/AgentSessionIntegrationTestSupport.kt @@ -452,6 +452,8 @@ internal suspend fun withTestServiceAndLaunch( toolWindowVisibleFlow: StateFlow = MutableStateFlow(true), currentTimeMillis: () -> Long = System::currentTimeMillis, loadingDelayMs: Long = 0L, + visibleCostHydrationDelayMs: Long = 750L, + workingThreadCostCacheTtlMs: Long = 60_000L, branchMismatchConfirmation: suspend (Project?, String, String) -> Boolean = { _, _, _ -> error("Unexpected branch mismatch confirmation") }, @@ -472,6 +474,8 @@ internal suspend fun withTestServiceAndLaunch( toolWindowVisibleFlow = toolWindowVisibleFlow, currentTimeMillis = currentTimeMillis, loadingDelayMs = loadingDelayMs, + visibleCostHydrationDelayMs = visibleCostHydrationDelayMs, + workingThreadCostCacheTtlMs = workingThreadCostCacheTtlMs, branchMismatchConfirmation = branchMismatchConfirmation, action = action, ) @@ -501,6 +505,8 @@ internal suspend fun withService( toolWindowVisibleFlow: StateFlow = MutableStateFlow(true), currentTimeMillis: () -> Long = System::currentTimeMillis, loadingDelayMs: Long = 0L, + visibleCostHydrationDelayMs: Long = 750L, + workingThreadCostCacheTtlMs: Long = 60_000L, action: suspend (AgentSessionStateSyncTestFacade) -> Unit, ) { withServiceAndLaunch( @@ -518,6 +524,8 @@ internal suspend fun withService( toolWindowVisibleFlow = toolWindowVisibleFlow, currentTimeMillis = currentTimeMillis, loadingDelayMs = loadingDelayMs, + visibleCostHydrationDelayMs = visibleCostHydrationDelayMs, + workingThreadCostCacheTtlMs = workingThreadCostCacheTtlMs, ) { service, _ -> action(service) } @@ -547,6 +555,8 @@ internal suspend fun withServiceAndLaunch( toolWindowVisibleFlow: StateFlow = MutableStateFlow(true), currentTimeMillis: () -> Long = System::currentTimeMillis, loadingDelayMs: Long = 0L, + visibleCostHydrationDelayMs: Long = 750L, + workingThreadCostCacheTtlMs: Long = 60_000L, branchMismatchConfirmation: suspend (Project?, String, String) -> Boolean = { _, _, _ -> error("Unexpected branch mismatch confirmation") }, @@ -568,6 +578,8 @@ internal suspend fun withServiceAndLaunch( toolWindowVisibleFlow = toolWindowVisibleFlow, currentTimeMillis = currentTimeMillis, loadingDelayMs = loadingDelayMs, + visibleCostHydrationDelayMs = visibleCostHydrationDelayMs, + workingThreadCostCacheTtlMs = workingThreadCostCacheTtlMs, branchMismatchConfirmation = branchMismatchConfirmation, ) { service, _, launchService -> action(service, launchService) @@ -600,6 +612,8 @@ internal suspend fun withServiceAndArchive( toolWindowVisibleFlow: StateFlow = MutableStateFlow(true), currentTimeMillis: () -> Long = System::currentTimeMillis, loadingDelayMs: Long = 0L, + visibleCostHydrationDelayMs: Long = 750L, + workingThreadCostCacheTtlMs: Long = 60_000L, action: suspend (AgentSessionStateSyncTestFacade, AgentSessionArchiveService) -> Unit, ) { withServiceAndArchiveAndLaunch( @@ -619,6 +633,8 @@ internal suspend fun withServiceAndArchive( toolWindowVisibleFlow = toolWindowVisibleFlow, currentTimeMillis = currentTimeMillis, loadingDelayMs = loadingDelayMs, + visibleCostHydrationDelayMs = visibleCostHydrationDelayMs, + workingThreadCostCacheTtlMs = workingThreadCostCacheTtlMs, ) { service, archiveService, _ -> action(service, archiveService) } @@ -650,6 +666,8 @@ internal suspend fun withServiceAndArchiveAndLaunch( toolWindowVisibleFlow: StateFlow = MutableStateFlow(true), currentTimeMillis: () -> Long = System::currentTimeMillis, loadingDelayMs: Long = 0L, + visibleCostHydrationDelayMs: Long = 750L, + workingThreadCostCacheTtlMs: Long = 60_000L, branchMismatchConfirmation: suspend (Project?, String, String) -> Boolean = { _, _, _ -> error("Unexpected branch mismatch confirmation") }, @@ -691,6 +709,8 @@ internal suspend fun withServiceAndArchiveAndLaunch( currentTimeMillis = currentTimeMillis, archiveTransitionSuppressions = archiveTransitionSuppressions, loadingDelayMs = loadingDelayMs, + visibleCostHydrationDelayMs = visibleCostHydrationDelayMs, + workingThreadCostCacheTtlMs = workingThreadCostCacheTtlMs, subscribeToProjectLifecycle = false, ) val service = AgentSessionStateSyncTestFacade( diff --git a/plugins/agent-workbench/sessions/testSrc/AgentSessionRefreshServiceIntegrationTest.kt b/plugins/agent-workbench/sessions/testSrc/AgentSessionRefreshServiceIntegrationTest.kt index 1f3cfc9575be..0856ace66d9c 100644 --- a/plugins/agent-workbench/sessions/testSrc/AgentSessionRefreshServiceIntegrationTest.kt +++ b/plugins/agent-workbench/sessions/testSrc/AgentSessionRefreshServiceIntegrationTest.kt @@ -189,6 +189,69 @@ class AgentSessionRefreshServiceIntegrationTest { } } + @Test + fun workingVisibleThreadCostReloadsAfterTtlWithoutStateRefresh() = runBlocking(Dispatchers.Default) { + var nowMs = 1_000L + val costLoadCount = AtomicInteger(0) + + withService( + sessionSourcesProvider = { + listOf( + ScriptedSessionSource( + provider = AgentSessionProvider.from("claude"), + listFromOpenProject = { _, _ -> + listOf( + thread( + id = "claude-1", + updatedAt = 100L, + activity = AgentThreadActivity.PROCESSING, + provider = AgentSessionProvider.from("claude"), + ) + ) + }, + loadThreadCostsProvider = { _, requestedThreads -> + val loadNumber = costLoadCount.incrementAndGet() + requestedThreads.associate { thread -> + thread.id to AgentSessionCost( + amountUsd = BigDecimal(loadNumber), + kind = AgentSessionCostKind.EXACT, + ) + } + }, + ) + ) + }, + projectEntriesProvider = { + listOf(openProjectEntry(PROJECT_PATH, "Project A")) + }, + currentTimeMillis = { nowMs }, + visibleCostHydrationDelayMs = 10L, + workingThreadCostCacheTtlMs = 100L, + ) { service -> + service.refresh() + + waitForCondition { + service.state.value.projects.firstOrNull { it.path == PROJECT_PATH } + ?.threads + ?.singleOrNull() + ?.cost + ?.amountUsd == BigDecimal.ONE + } + assertThat(costLoadCount.get()).isEqualTo(1) + + nowMs += 101L + + waitForCondition { + service.state.value.projects.firstOrNull { it.path == PROJECT_PATH } + ?.threads + ?.singleOrNull() + ?.cost + ?.amountUsd == BigDecimal.valueOf(2) + } + assertThat(costLoadCount.get()).isEqualTo(2) + } + } + @Test fun persistedWarmSnapshotCostSurvivesRefreshAndSkipsReload() = runBlocking(Dispatchers.Default) { val persistedWarmState = AgentSessionWarmStateService()