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

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 {

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<AgentSessionVisibleCostHydrationSupport>()
-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<AgentSessionSource>,
   private val toolWindowVisibleFlow: StateFlow<Boolean>,
   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<ThreadCacheKey, ThreadCostCacheEntry>()
   private val inFlightLoads = ConcurrentHashMap.newKeySet<ThreadLoadKey>()
@@ -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(

plugins/agent-workbench/sessions/testSrc/AgentSessionIntegrationTestSupport.kt

@@ -452,6 +452,8 @@ internal suspend fun withTestServiceAndLaunch(
   toolWindowVisibleFlow: StateFlow<Boolean> = 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<Boolean> = 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<Boolean> = 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<Boolean> = 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<Boolean> = 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(

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()