IJPL-233558 split AgentSessionsService

GitOrigin-RevId: 88f04faad31c8140f8c9ba96e2b83aca1a6183dc
2026-09-27 10:03:11 +07:00 · 2026-02-16 19:03:38 +00:00
parent 15ca8a4acf
commit 174490c8ed
7 changed files with 1274 additions and 1122 deletions
@@ -0,0 +1,697 @@
+// Copyright 2000-2026 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.agent.workbench.sessions
+
+import com.intellij.agent.workbench.sessions.providers.AgentSessionProviderBridges
+import com.intellij.agent.workbench.sessions.providers.AgentSessionSource
+import com.intellij.openapi.diagnostic.logger
+import com.intellij.openapi.project.Project
+import kotlinx.coroutines.CancellationException
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.async
+import kotlinx.coroutines.awaitAll
+import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
+import kotlin.time.Duration.Companion.milliseconds
+
+private val LOG = logger<AgentSessionsLoadingCoordinator>()
+private const val SOURCE_UPDATE_DEBOUNCE_MS = 350L
+
+internal class AgentSessionsLoadingCoordinator(
+  private val serviceScope: CoroutineScope,
+  private val sessionSourcesProvider: () -> List<AgentSessionSource>,
+  private val projectEntriesProvider: suspend () -> List<ProjectEntry>,
+  private val treeUiState: SessionsTreeUiState,
+  private val stateStore: AgentSessionsStateStore,
+  private val isRefreshGateActive: suspend () -> Boolean,
+) {
+  private val refreshMutex = Mutex()
+  private val onDemandMutex = Mutex()
+  private val onDemandLoading = LinkedHashSet<String>()
+  private val onDemandWorktreeLoading = LinkedHashSet<String>()
+  private val sourceRefreshJobs = LinkedHashMap<AgentSessionProvider, Job>()
+  private val sourceRefreshJobsLock = Any()
+
+  fun observeSessionSourceUpdates() {
+    serviceScope.launch {
+      for (source in sessionSourcesProvider()) {
+        launch {
+          source.updates.collect {
+            scheduleSourceRefresh(source.provider)
+          }
+        }
+      }
+    }
+  }
+
+  fun refresh() {
+    serviceScope.launch(Dispatchers.IO) {
+      if (!refreshMutex.tryLock()) {
+        return@launch
+      }
+      try {
+        val entries = projectEntriesProvider()
+        val currentState = stateStore.snapshot()
+        val currentProjectsByPath = currentState.projects.associateBy { normalizePath(it.path) }
+        val openPaths = entries.flatMap { entry ->
+          buildList {
+            if (entry.project != null) add(normalizePath(entry.path))
+            entry.worktreeEntries.filter { it.project != null }.forEach { add(normalizePath(it.path)) }
+          }
+        }
+        treeUiState.retainOpenProjectThreadPreviews(openPaths.toSet())
+        val knownPaths = entries.flatMap { entry ->
+          buildList {
+            add(normalizePath(entry.path))
+            entry.worktreeEntries.forEach { add(normalizePath(it.path)) }
+          }
+        }
+        val initialVisibleThreadCounts = stateStore.buildInitialVisibleThreadCounts(knownPaths)
+        val initialProjects = entries.map { entry ->
+          val normalizedEntryPath = normalizePath(entry.path)
+          val existing = currentProjectsByPath[normalizedEntryPath]
+          val cachedPreviews = if (entry.project != null) {
+            treeUiState.getOpenProjectThreadPreviews(normalizedEntryPath)
+          }
+          else {
+            null
+          }
+          val cachedThreads = cachedPreviews.orEmpty().toCachedSessionThreads()
+          AgentProjectSessions(
+            path = normalizedEntryPath,
+            name = entry.name,
+            branch = entry.branch ?: existing?.branch,
+            isOpen = entry.project != null,
+            isLoading = entry.project != null,
+            hasLoaded = existing?.hasLoaded ?: (cachedPreviews != null),
+            hasUnknownThreadCount = existing?.hasUnknownThreadCount ?: false,
+            threads = existing?.threads ?: cachedThreads,
+            errorMessage = existing?.errorMessage,
+            providerWarnings = existing?.providerWarnings ?: emptyList(),
+            worktrees = entry.worktreeEntries.map { wt ->
+              val normalizedWorktreePath = normalizePath(wt.path)
+              val existingWt = existing?.worktrees?.firstOrNull { normalizePath(it.path) == normalizedWorktreePath }
+              val cachedWorktreePreviews = if (wt.project != null) {
+                treeUiState.getOpenProjectThreadPreviews(normalizedWorktreePath)
+              }
+              else {
+                null
+              }
+              val cachedWorktreeThreads = cachedWorktreePreviews.orEmpty().toCachedSessionThreads()
+              val hasExistingData = existingWt != null && existingWt.threads.isNotEmpty()
+              AgentWorktree(
+                path = normalizedWorktreePath,
+                name = wt.name,
+                branch = wt.branch,
+                isOpen = wt.project != null,
+                isLoading = wt.project != null && (hasExistingData || cachedWorktreePreviews != null),
+                hasLoaded = existingWt?.hasLoaded ?: (cachedWorktreePreviews != null),
+                hasUnknownThreadCount = existingWt?.hasUnknownThreadCount ?: false,
+                threads = existingWt?.threads ?: cachedWorktreeThreads,
+                errorMessage = existingWt?.errorMessage,
+                providerWarnings = existingWt?.providerWarnings ?: emptyList(),
+              )
+            },
+          )
+        }
+        stateStore.replaceProjects(
+          projects = initialProjects,
+          visibleThreadCounts = initialVisibleThreadCounts,
+        )
+
+        val sessionSources = sessionSourcesProvider()
+
+        // Prefetch from all sources in parallel
+        val prefetchedByProvider = coroutineScope {
+          sessionSources.map { source ->
+            async {
+              source.provider to try {
+                source.prefetchThreads(openPaths)
+              }
+              catch (_: Throwable) {
+                emptyMap()
+              }
+            }
+          }.awaitAll().toMap()
+        }
+
+        // Load each (project × source) independently so fast sources (Claude)
+        // update the UI immediately without waiting for slow sources (Codex).
+        coroutineScope {
+          for (entry in entries) {
+            launch {
+              val normalizedEntryPath = normalizePath(entry.path)
+              if (entry.project == null) {
+                stateStore.updateProject(normalizedEntryPath) { it.copy(isLoading = false) }
+                return@launch
+              }
+              val sourceResults = java.util.concurrent.CopyOnWriteArrayList<AgentSessionSourceLoadResult>()
+              coroutineScope {
+                for (source in sessionSources) {
+                  launch {
+                    val sourceResult = loadSourceResultForOpenProject(
+                      source = source,
+                      normalizedPath = normalizedEntryPath,
+                      project = entry.project,
+                      prefetchedByProvider = prefetchedByProvider,
+                      originalPath = entry.path,
+                    )
+                    sourceResults.add(sourceResult)
+                    // Incremental UI update — clear spinner as soon as any source succeeds
+                    val partial = mergeAgentSessionSourceLoadResults(
+                      sourceResults = sourceResults.toList(),
+                      resolveErrorMessage = ::resolveErrorMessage,
+                      resolveWarningMessage = ::resolveProviderWarningMessage,
+                    )
+                    val anySuccess = sourceResults.any { it.result.isSuccess }
+                    stateStore.updateProject(normalizedEntryPath) { project ->
+                      project.copy(
+                        threads = partial.threads,
+                        providerWarnings = partial.providerWarnings,
+                        isLoading = if (anySuccess) false else project.isLoading,
+                      )
+                    }
+                  }
+                }
+              }
+              // All sources done — final update with error/warning consolidation
+              val finalResult = mergeAgentSessionSourceLoadResults(
+                sourceResults = sourceResults.toList(),
+                resolveErrorMessage = ::resolveErrorMessage,
+                resolveWarningMessage = ::resolveProviderWarningMessage,
+              )
+              stateStore.updateProject(normalizedEntryPath) { project ->
+                project.copy(
+                  isLoading = false,
+                  hasLoaded = true,
+                  hasUnknownThreadCount = finalResult.hasUnknownThreadCount,
+                  threads = finalResult.threads,
+                  errorMessage = finalResult.errorMessage,
+                  providerWarnings = finalResult.providerWarnings,
+                )
+              }
+              if (finalResult.errorMessage == null) {
+                treeUiState.setOpenProjectThreadPreviews(normalizedEntryPath, finalResult.threads.toThreadPreviews())
+              }
+            }
+            for (wt in entry.worktreeEntries) {
+              launch {
+                val normalizedEntryPath = normalizePath(entry.path)
+                val normalizedWorktreePath = normalizePath(wt.path)
+                if (wt.project == null) {
+                  stateStore.updateWorktree(normalizedEntryPath, normalizedWorktreePath) { it.copy(isLoading = false) }
+                  return@launch
+                }
+                val sourceResults = java.util.concurrent.CopyOnWriteArrayList<AgentSessionSourceLoadResult>()
+                coroutineScope {
+                  for (source in sessionSources) {
+                    launch {
+                      val sourceResult = loadSourceResultForOpenProject(
+                        source = source,
+                        normalizedPath = normalizedWorktreePath,
+                        project = wt.project,
+                        prefetchedByProvider = prefetchedByProvider,
+                        originalPath = wt.path,
+                      )
+                      sourceResults.add(sourceResult)
+                      val partial = mergeAgentSessionSourceLoadResults(
+                        sourceResults = sourceResults.toList(),
+                        resolveErrorMessage = ::resolveErrorMessage,
+                        resolveWarningMessage = ::resolveProviderWarningMessage,
+                      )
+                      val anySuccess = sourceResults.any { it.result.isSuccess }
+                      stateStore.updateWorktree(normalizedEntryPath, normalizedWorktreePath) { worktree ->
+                        worktree.copy(
+                          threads = partial.threads,
+                          providerWarnings = partial.providerWarnings,
+                          isLoading = if (anySuccess) false else worktree.isLoading,
+                        )
+                      }
+                    }
+                  }
+                }
+                val finalResult = mergeAgentSessionSourceLoadResults(
+                  sourceResults = sourceResults.toList(),
+                  resolveErrorMessage = ::resolveErrorMessage,
+                  resolveWarningMessage = ::resolveProviderWarningMessage,
+                )
+                stateStore.updateWorktree(normalizedEntryPath, normalizedWorktreePath) { worktree ->
+                  worktree.copy(
+                    isLoading = false,
+                    hasLoaded = true,
+                    hasUnknownThreadCount = finalResult.hasUnknownThreadCount,
+                    threads = finalResult.threads,
+                    errorMessage = finalResult.errorMessage,
+                    providerWarnings = finalResult.providerWarnings,
+                  )
+                }
+                if (finalResult.errorMessage == null) {
+                  treeUiState.setOpenProjectThreadPreviews(normalizedWorktreePath, finalResult.threads.toThreadPreviews())
+                }
+              }
+            }
+          }
+        }
+        stateStore.update { it.copy(lastUpdatedAt = System.currentTimeMillis()) }
+      }
+      catch (e: Throwable) {
+        if (e is CancellationException) throw e
+        LOG.error("Failed to load agent sessions", e)
+        stateStore.markLoadFailure(AgentSessionsBundle.message("toolwindow.error"))
+      }
+      finally {
+        refreshMutex.unlock()
+      }
+    }
+  }
+
+  fun loadProjectThreadsOnDemand(path: String) {
+    serviceScope.launch(Dispatchers.IO) {
+      val normalized = normalizePath(path)
+      if (!markOnDemandLoading(normalized)) return@launch
+      try {
+        stateStore.updateProject(normalized) { project ->
+          project.copy(
+            isLoading = true,
+            hasUnknownThreadCount = false,
+            errorMessage = null,
+            providerWarnings = emptyList(),
+          )
+        }
+        val result = loadThreadsFromClosedProject(path = normalized)
+        stateStore.updateProject(normalized) { project ->
+          project.copy(
+            isLoading = false,
+            hasLoaded = true,
+            hasUnknownThreadCount = result.hasUnknownThreadCount,
+            threads = result.threads,
+            errorMessage = result.errorMessage,
+            providerWarnings = result.providerWarnings,
+          )
+        }
+      }
+      finally {
+        clearOnDemandLoading(normalized)
+      }
+    }
+  }
+
+  fun loadWorktreeThreadsOnDemand(projectPath: String, worktreePath: String) {
+    serviceScope.launch(Dispatchers.IO) {
+      val normalizedProject = normalizePath(projectPath)
+      val normalizedWorktree = normalizePath(worktreePath)
+      if (!markWorktreeOnDemandLoading(normalizedProject, normalizedWorktree)) return@launch
+      try {
+        stateStore.updateWorktree(normalizedProject, normalizedWorktree) { worktree ->
+          worktree.copy(
+            isLoading = true,
+            hasUnknownThreadCount = false,
+            errorMessage = null,
+            providerWarnings = emptyList(),
+          )
+        }
+        val result = loadThreadsFromClosedProject(path = normalizedWorktree)
+        stateStore.updateWorktree(normalizedProject, normalizedWorktree) { worktree ->
+          worktree.copy(
+            isLoading = false,
+            hasLoaded = true,
+            hasUnknownThreadCount = result.hasUnknownThreadCount,
+            threads = result.threads,
+            errorMessage = result.errorMessage,
+            providerWarnings = result.providerWarnings,
+          )
+        }
+      }
+      finally {
+        clearWorktreeOnDemandLoading(normalizedWorktree)
+      }
+    }
+  }
+
+  fun appendProviderUnavailableWarning(path: String, provider: AgentSessionProvider) {
+    val warning = AgentSessionProviderWarning(provider = provider, message = providerUnavailableMessage(provider))
+    stateStore.update { state ->
+      var updated = false
+      val nextProjects = state.projects.map { project ->
+        if (project.path == path) {
+          updated = true
+          project.copy(providerWarnings = mergeProviderWarning(project.providerWarnings, warning))
+        }
+        else {
+          val nextWorktrees = project.worktrees.map { worktree ->
+            if (worktree.path == path) {
+              updated = true
+              worktree.copy(providerWarnings = mergeProviderWarning(worktree.providerWarnings, warning))
+            }
+            else {
+              worktree
+            }
+          }
+          if (nextWorktrees == project.worktrees) project else project.copy(worktrees = nextWorktrees)
+        }
+      }
+      if (!updated) state else state.copy(projects = nextProjects, lastUpdatedAt = System.currentTimeMillis())
+    }
+  }
+
+  private fun scheduleSourceRefresh(provider: AgentSessionProvider) {
+    synchronized(sourceRefreshJobsLock) {
+      sourceRefreshJobs.remove(provider)?.cancel()
+      val job = serviceScope.launch(Dispatchers.IO) {
+        delay(SOURCE_UPDATE_DEBOUNCE_MS.milliseconds)
+        if (!isRefreshGateActive()) return@launch
+        refreshLoadedProviderThreads(provider)
+      }
+      sourceRefreshJobs[provider] = job
+      job.invokeOnCompletion {
+        synchronized(sourceRefreshJobsLock) {
+          if (sourceRefreshJobs[provider] === job) {
+            sourceRefreshJobs.remove(provider)
+          }
+        }
+      }
+    }
+  }
+
+  private suspend fun refreshLoadedProviderThreads(provider: AgentSessionProvider) {
+    if (!refreshMutex.tryLock()) return
+    try {
+      val source = sessionSourcesProvider().firstOrNull { it.provider == provider } ?: return
+      val stateSnapshot = stateStore.snapshot()
+      val targetPaths = collectLoadedPaths(stateSnapshot)
+      if (targetPaths.isEmpty()) return
+
+      val prefetched = try {
+        source.prefetchThreads(targetPaths)
+      }
+      catch (_: Throwable) {
+        emptyMap()
+      }
+
+      val outcomes = LinkedHashMap<String, ProviderRefreshOutcome>(targetPaths.size)
+      for (path in targetPaths) {
+        val prefetchedThreads = prefetched[path]
+        if (prefetchedThreads != null) {
+          outcomes[path] = ProviderRefreshOutcome(threads = prefetchedThreads)
+          continue
+        }
+
+        try {
+          outcomes[path] = ProviderRefreshOutcome(threads = source.listThreadsFromClosedProject(path))
+        }
+        catch (e: Throwable) {
+          if (e is CancellationException) throw e
+          LOG.warn("Failed to refresh ${provider.value} sessions for $path", e)
+          outcomes[path] = ProviderRefreshOutcome(
+            warningMessage = resolveProviderWarningMessage(provider, e),
+          )
+        }
+      }
+
+      stateStore.update { state ->
+        var changed = false
+        val nextProjects = state.projects.map { project ->
+          val updatedProject = if (project.hasLoaded) {
+            val outcome = outcomes[project.path]
+            if (outcome != null) {
+              changed = true
+              project.withProviderRefreshOutcome(provider, outcome)
+            }
+            else {
+              project
+            }
+          }
+          else {
+            project
+          }
+
+          val nextWorktrees = updatedProject.worktrees.map { worktree ->
+            if (!worktree.hasLoaded) return@map worktree
+            val outcome = outcomes[worktree.path] ?: return@map worktree
+            changed = true
+            worktree.withProviderRefreshOutcome(provider, outcome)
+          }
+
+          if (nextWorktrees == updatedProject.worktrees) {
+            updatedProject
+          }
+          else {
+            updatedProject.copy(worktrees = nextWorktrees)
+          }
+        }
+
+        if (!changed) {
+          state
+        }
+        else {
+          state.copy(
+            projects = nextProjects,
+            lastUpdatedAt = System.currentTimeMillis(),
+          )
+        }
+      }
+    }
+    finally {
+      refreshMutex.unlock()
+    }
+  }
+
+  private fun collectLoadedPaths(state: AgentSessionsState): List<String> {
+    val paths = LinkedHashSet<String>()
+    for (project in state.projects) {
+      if (project.hasLoaded) {
+        paths.add(project.path)
+      }
+      for (worktree in project.worktrees) {
+        if (worktree.hasLoaded) {
+          paths.add(worktree.path)
+        }
+      }
+    }
+    return ArrayList(paths)
+  }
+
+  private suspend fun markOnDemandLoading(path: String): Boolean {
+    return onDemandMutex.withLock {
+      val project = stateStore.state.value.projects.firstOrNull { it.path == path } ?: return@withLock false
+      if (project.isOpen || project.isLoading || project.hasLoaded) return@withLock false
+      if (!onDemandLoading.add(path)) return@withLock false
+      true
+    }
+  }
+
+  private suspend fun clearOnDemandLoading(path: String) {
+    onDemandMutex.withLock {
+      onDemandLoading.remove(path)
+    }
+  }
+
+  private suspend fun markWorktreeOnDemandLoading(projectPath: String, worktreePath: String): Boolean {
+    return onDemandMutex.withLock {
+      val project = stateStore.state.value.projects.firstOrNull { it.path == projectPath } ?: return@withLock false
+      val worktree = project.worktrees.firstOrNull { it.path == worktreePath } ?: return@withLock false
+      if (worktree.isLoading || worktree.hasLoaded) return@withLock false
+      if (!onDemandWorktreeLoading.add(worktreePath)) return@withLock false
+      true
+    }
+  }
+
+  private suspend fun clearWorktreeOnDemandLoading(worktreePath: String) {
+    onDemandMutex.withLock {
+      onDemandWorktreeLoading.remove(worktreePath)
+    }
+  }
+
+  private suspend fun loadThreadsFromClosedProject(path: String): AgentSessionLoadResult {
+    return loadThreads(path) { source ->
+      source.listThreadsFromClosedProject(path = path)
+    }
+  }
+
+  private suspend fun loadThreads(
+    path: String,
+    loadOperation: suspend (AgentSessionSource) -> List<AgentSessionThread>,
+  ): AgentSessionLoadResult {
+    val sessionSources = sessionSourcesProvider()
+    val sourceResults = coroutineScope {
+      sessionSources.map { source ->
+        async {
+          val result = try {
+            Result.success(loadOperation(source))
+          }
+          catch (throwable: Throwable) {
+            if (throwable is CancellationException) throw throwable
+            LOG.warn("Failed to load ${source.provider.value} sessions for $path", throwable)
+            Result.failure(throwable)
+          }
+          AgentSessionSourceLoadResult(
+            provider = source.provider,
+            result = result,
+            hasUnknownTotal = result.isSuccess && !source.canReportExactThreadCount,
+          )
+        }
+      }.awaitAll()
+    }
+    return mergeAgentSessionSourceLoadResults(
+      sourceResults = sourceResults,
+      resolveErrorMessage = ::resolveErrorMessage,
+      resolveWarningMessage = ::resolveProviderWarningMessage,
+    )
+  }
+
+  private suspend fun loadSourceResultForOpenProject(
+    source: AgentSessionSource,
+    normalizedPath: String,
+    project: Project,
+    prefetchedByProvider: Map<AgentSessionProvider, Map<String, List<AgentSessionThread>>>,
+    originalPath: String,
+  ): AgentSessionSourceLoadResult {
+    return try {
+      val prefetched = prefetchedByProvider[source.provider]?.get(normalizedPath)
+      val threads = prefetched ?: source.listThreadsFromOpenProject(path = normalizedPath, project = project)
+      AgentSessionSourceLoadResult(
+        provider = source.provider,
+        result = Result.success(threads),
+        hasUnknownTotal = !source.canReportExactThreadCount,
+      )
+    }
+    catch (e: Throwable) {
+      if (e is CancellationException) throw e
+      LOG.warn("Failed to load ${source.provider.value} sessions for $originalPath", e)
+      AgentSessionSourceLoadResult(
+        provider = source.provider,
+        result = Result.failure(e),
+      )
+    }
+  }
+
+  private fun resolveErrorMessage(provider: AgentSessionProvider, t: Throwable): String {
+    return if (isCliMissingError(provider, t)) resolveCliMissingMessage(provider)
+    else AgentSessionsBundle.message("toolwindow.error")
+  }
+
+  private fun resolveCliMissingMessage(provider: AgentSessionProvider): String {
+    return if (AgentSessionProviderBridges.find(provider) != null) {
+      AgentSessionsBundle.message(agentSessionCliMissingMessageKey(provider))
+    }
+    else {
+      providerUnavailableMessage(provider)
+    }
+  }
+
+  private fun resolveProviderWarningMessage(provider: AgentSessionProvider, t: Throwable): String {
+    return if (isCliMissingError(provider, t)) resolveCliMissingMessage(provider)
+    else AgentSessionsBundle.message("toolwindow.warning.provider.unavailable", resolveProviderLabel(provider))
+  }
+
+  private fun isCliMissingError(provider: AgentSessionProvider, t: Throwable): Boolean {
+    return AgentSessionProviderBridges.find(provider)?.isCliMissingError(t) == true
+  }
+
+  private fun resolveProviderLabel(provider: AgentSessionProvider): String {
+    val bridge = AgentSessionProviderBridges.find(provider)
+    return if (bridge != null) AgentSessionsBundle.message(bridge.displayNameKey) else provider.value
+  }
+
+  private fun providerUnavailableMessage(provider: AgentSessionProvider): String {
+    return AgentSessionsBundle.message("toolwindow.warning.provider.unavailable", resolveProviderLabel(provider))
+  }
+
+  private fun mergeProviderWarning(
+    warnings: List<AgentSessionProviderWarning>,
+    warning: AgentSessionProviderWarning,
+  ): List<AgentSessionProviderWarning> {
+    if (warnings.any { it.provider == warning.provider && it.message == warning.message }) {
+      return warnings
+    }
+    return warnings + warning
+  }
+
+  private fun AgentProjectSessions.withProviderRefreshOutcome(
+    provider: AgentSessionProvider,
+    outcome: ProviderRefreshOutcome,
+  ): AgentProjectSessions {
+    val mergedThreads = outcome.threads?.let { threads ->
+      mergeThreadsForProvider(this.threads, provider, threads)
+    } ?: this.threads
+    return copy(
+      threads = mergedThreads,
+      providerWarnings = replaceProviderWarning(this.providerWarnings, provider, outcome.warningMessage),
+    )
+  }
+
+  private fun AgentWorktree.withProviderRefreshOutcome(
+    provider: AgentSessionProvider,
+    outcome: ProviderRefreshOutcome,
+  ): AgentWorktree {
+    val mergedThreads = outcome.threads?.let { threads ->
+      mergeThreadsForProvider(this.threads, provider, threads)
+    } ?: this.threads
+    return copy(
+      threads = mergedThreads,
+      providerWarnings = replaceProviderWarning(this.providerWarnings, provider, outcome.warningMessage),
+    )
+  }
+
+  private fun replaceProviderWarning(
+    warnings: List<AgentSessionProviderWarning>,
+    provider: AgentSessionProvider,
+    warningMessage: String?,
+  ): List<AgentSessionProviderWarning> {
+    val withoutProvider = warnings.filterNot { it.provider == provider }
+    return if (warningMessage == null) {
+      withoutProvider
+    }
+    else {
+      withoutProvider + AgentSessionProviderWarning(provider = provider, message = warningMessage)
+    }
+  }
+
+  private fun mergeThreadsForProvider(
+    existingThreads: List<AgentSessionThread>,
+    provider: AgentSessionProvider,
+    newProviderThreads: List<AgentSessionThread>,
+  ): List<AgentSessionThread> {
+    val mergedThreads = ArrayList<AgentSessionThread>(existingThreads.size + newProviderThreads.size)
+    existingThreads.filterTo(mergedThreads) { it.provider != provider }
+    mergedThreads.addAll(newProviderThreads)
+    mergedThreads.sortByDescending { it.updatedAt }
+    return mergedThreads
+  }
+
+  private fun List<AgentSessionThreadPreview>.toCachedSessionThreads(): List<AgentSessionThread> {
+    return map { preview ->
+      AgentSessionThread(
+        id = preview.id,
+        title = preview.title,
+        updatedAt = preview.updatedAt,
+        archived = false,
+        provider = preview.provider,
+      )
+    }
+  }
+
+  private fun List<AgentSessionThread>.toThreadPreviews(): List<AgentSessionThreadPreview> {
+    return map { thread ->
+      AgentSessionThreadPreview(
+        id = thread.id,
+        title = thread.title,
+        updatedAt = thread.updatedAt,
+        provider = thread.provider,
+      )
+    }
+  }
+
+  private fun normalizePath(path: String): String {
+    return normalizeSessionsProjectPath(path)
+  }
+
+  private data class ProviderRefreshOutcome(
+    val threads: List<AgentSessionThread>? = null,
+    val warningMessage: String? = null,
+  )
+}
@@ -0,0 +1,195 @@
+// Copyright 2000-2026 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.agent.workbench.sessions
+
+import com.intellij.ide.RecentProjectsManager
+import com.intellij.ide.RecentProjectsManagerBase
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.project.ProjectManager
+import com.intellij.openapi.util.io.FileUtilRt
+import kotlinx.coroutines.async
+import kotlinx.coroutines.awaitAll
+import kotlinx.coroutines.coroutineScope
+import java.nio.file.InvalidPathException
+import java.nio.file.Path
+import kotlin.io.path.invariantSeparatorsPathString
+import kotlin.io.path.name
+
+internal class AgentSessionsProjectCatalog {
+  suspend fun collectProjects(): List<ProjectEntry> {
+    val rawEntries = collectRawProjectEntries()
+    if (rawEntries.isEmpty()) return emptyList()
+
+    val repoRootByPath = rawEntries.associate { entry ->
+      entry.path to GitWorktreeDiscovery.detectRepoRoot(entry.path)
+    }
+
+    data class RepoGroup(
+      val repoRoot: String,
+      val members: MutableList<IndexedValue<ProjectEntry>>,
+    )
+
+    val repoGroups = LinkedHashMap<String, RepoGroup>()
+    val standaloneEntries = mutableListOf<IndexedValue<ProjectEntry>>()
+
+    rawEntries.forEachIndexed { index, entry ->
+      val repoRoot = repoRootByPath[entry.path]
+      if (repoRoot != null) {
+        val group = repoGroups.getOrPut(repoRoot) {
+          RepoGroup(repoRoot, mutableListOf())
+        }
+        group.members.add(IndexedValue(index, entry))
+      }
+      else {
+        standaloneEntries.add(IndexedValue(index, entry))
+      }
+    }
+
+    // Discover all worktrees in parallel across repo roots (main + linked).
+    val discoveredByRepoRoot = coroutineScope {
+      repoGroups.keys.map { repoRoot ->
+        async { repoRoot to GitWorktreeDiscovery.discoverWorktrees(repoRoot) }
+      }.awaitAll().toMap()
+    }
+
+    val resultEntries = mutableListOf<IndexedValue<ProjectEntry>>()
+    for ((repoRoot, group) in repoGroups) {
+      val mainRaw = group.members.firstOrNull { it.value.path == repoRoot }
+      val worktreeRaws = group.members.filter { it.value.path != repoRoot }
+      val firstIndex = group.members.minOf { it.index }
+
+      val discoveredWorktrees = discoveredByRepoRoot[repoRoot] ?: emptyList()
+      val worktreeEntries = buildWorktreeEntries(worktreeRaws.map { it.value }, discoveredWorktrees)
+      val mainBranch = shortBranchName(discoveredWorktrees.firstOrNull { it.isMain }?.branch)
+
+      if (worktreeEntries.isEmpty()) {
+        val raw = mainRaw?.value ?: continue
+        resultEntries.add(IndexedValue(firstIndex, raw))
+      }
+      else {
+        val entry = mainRaw?.value?.copy(worktreeEntries = worktreeEntries, branch = mainBranch)
+                    ?: ProjectEntry(
+                      path = repoRoot,
+                      name = worktreeDisplayName(repoRoot),
+                      project = null,
+                      branch = mainBranch,
+                      worktreeEntries = worktreeEntries,
+                    )
+        resultEntries.add(IndexedValue(firstIndex, entry))
+      }
+    }
+
+    for (indexed in standaloneEntries) {
+      resultEntries.add(indexed)
+    }
+
+    return resultEntries.sortedBy { it.index }.map { it.value }
+  }
+
+  private fun buildWorktreeEntries(
+    openRawEntries: List<ProjectEntry>,
+    discovered: List<GitWorktreeInfo>,
+  ): List<WorktreeEntry> {
+    val openPaths = openRawEntries.mapTo(LinkedHashSet()) { it.path }
+    val result = mutableListOf<WorktreeEntry>()
+
+    for (raw in openRawEntries) {
+      val gitInfo = discovered.firstOrNull { it.path == raw.path }
+      result.add(
+        WorktreeEntry(
+          path = raw.path,
+          name = raw.name,
+          branch = shortBranchName(gitInfo?.branch),
+          project = raw.project,
+        ),
+      )
+    }
+
+    for (info in discovered) {
+      if (info.path !in openPaths && !info.isMain) {
+        result.add(
+          WorktreeEntry(
+            path = info.path,
+            name = worktreeDisplayName(info.path),
+            branch = shortBranchName(info.branch),
+            project = null,
+          ),
+        )
+      }
+    }
+
+    return result
+  }
+
+  private fun collectRawProjectEntries(): List<ProjectEntry> {
+    val manager = RecentProjectsManager.getInstance() as? RecentProjectsManagerBase
+                  ?: return emptyList()
+    val dedicatedProjectPath = AgentWorkbenchDedicatedFrameProjectManager.dedicatedProjectPath()
+    val openProjects = ProjectManager.getInstance().openProjects
+    val openByPath = LinkedHashMap<String, Project>()
+    for (project in openProjects) {
+      val path = manager.getProjectPath(project)?.invariantSeparatorsPathString
+                 ?: project.basePath?.let(::normalizePath)
+                 ?: continue
+      if (path == dedicatedProjectPath || AgentWorkbenchDedicatedFrameProjectManager.isDedicatedProjectPath(path)) continue
+      openByPath[path] = project
+    }
+    val seen = LinkedHashSet<String>()
+    val entries = mutableListOf<ProjectEntry>()
+    for (path in manager.getRecentPaths()) {
+      val normalized = normalizePath(path)
+      if (normalized == dedicatedProjectPath || AgentWorkbenchDedicatedFrameProjectManager.isDedicatedProjectPath(normalized)) continue
+      if (!seen.add(normalized)) continue
+      entries.add(
+        ProjectEntry(
+          path = normalized,
+          name = resolveProjectName(manager, normalized, openByPath[normalized]),
+          project = openByPath[normalized],
+        ),
+      )
+    }
+    for ((path, project) in openByPath) {
+      if (!seen.add(path)) continue
+      entries.add(
+        ProjectEntry(
+          path = path,
+          name = resolveProjectName(manager, path, project),
+          project = project,
+        ),
+      )
+    }
+    return entries
+  }
+
+  private fun resolveProjectName(
+    manager: RecentProjectsManagerBase,
+    path: String,
+    project: Project?,
+  ): String {
+    val displayName = manager.getDisplayName(path).takeIf { !it.isNullOrBlank() }
+    if (displayName != null) return displayName
+    val projectName = manager.getProjectName(path)
+    if (projectName.isNotBlank()) return projectName
+    if (project != null) return project.name
+    return resolveProjectNameWithoutManager(path, project)
+  }
+
+  private fun resolveProjectNameWithoutManager(path: String, project: Project?): String {
+    if (project != null) return project.name
+    val fileName = try {
+      Path.of(path).name
+    }
+    catch (_: InvalidPathException) {
+      null
+    }
+    return fileName ?: FileUtilRt.toSystemDependentName(path)
+  }
+
+  private fun normalizePath(path: String): String {
+    return try {
+      Path.of(path).invariantSeparatorsPathString
+    }
+    catch (_: InvalidPathException) {
+      path
+    }
+  }
+}
@@ -0,0 +1,19 @@
+// Copyright 2000-2026 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.agent.workbench.sessions
+
+import com.intellij.openapi.project.Project
+
+internal data class ProjectEntry(
+  val path: String,
+  val name: String,
+  val project: Project?,
+  val branch: String? = null,
+  val worktreeEntries: List<WorktreeEntry> = emptyList(),
+)
+
+internal data class WorktreeEntry(
+  val path: String,
+  val name: String,
+  val branch: String?,
+  val project: Project?,
+)
@@ -0,0 +1,191 @@
+// Copyright 2000-2026 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.agent.workbench.sessions
+
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.update
+import java.nio.file.InvalidPathException
+import java.nio.file.Path
+import kotlin.io.path.invariantSeparatorsPathString
+
+internal class AgentSessionsStateStore(
+  private val treeUiState: SessionsTreeUiState,
+) {
+  private val mutableState = MutableStateFlow(AgentSessionsState())
+  val state: StateFlow<AgentSessionsState> = mutableState.asStateFlow()
+
+  fun snapshot(): AgentSessionsState = mutableState.value
+
+  fun update(transform: (AgentSessionsState) -> AgentSessionsState) {
+    mutableState.update(transform)
+  }
+
+  fun replaceProjects(projects: List<AgentProjectSessions>, visibleThreadCounts: Map<String, Int>) {
+    mutableState.update {
+      it.copy(
+        projects = projects,
+        visibleThreadCounts = visibleThreadCounts,
+        lastUpdatedAt = System.currentTimeMillis(),
+      )
+    }
+  }
+
+  fun markLoadFailure(errorMessage: String) {
+    mutableState.update { state ->
+      state.copy(
+        projects = state.projects.map { project ->
+          project.copy(
+            isLoading = false,
+            hasLoaded = true,
+            hasUnknownThreadCount = false,
+            errorMessage = errorMessage,
+            providerWarnings = emptyList(),
+            worktrees = project.worktrees.map { wt ->
+              wt.copy(isLoading = false, hasUnknownThreadCount = false, providerWarnings = emptyList())
+            },
+          )
+        },
+        lastUpdatedAt = System.currentTimeMillis(),
+      )
+    }
+  }
+
+  fun showMoreProjects() {
+    mutableState.update { it.copy(visibleProjectCount = it.visibleProjectCount + DEFAULT_VISIBLE_PROJECT_COUNT) }
+  }
+
+  fun showMoreThreads(path: String) {
+    val normalizedPath = normalizePath(path)
+    var deltaToPersist = 0
+    mutableState.update { state ->
+      val current = state.visibleThreadCounts[normalizedPath] ?: treeUiState.getVisibleThreadCount(normalizedPath)
+      val nextVisible = current + DEFAULT_VISIBLE_THREAD_COUNT
+      deltaToPersist = nextVisible - current
+      state.copy(visibleThreadCounts = state.visibleThreadCounts + (normalizedPath to nextVisible))
+    }
+    if (deltaToPersist > 0) {
+      treeUiState.incrementVisibleThreadCount(normalizedPath, deltaToPersist)
+    }
+  }
+
+  fun ensureThreadVisible(path: String, provider: AgentSessionProvider, threadId: String) {
+    val normalizedPath = normalizePath(path)
+    var deltaToPersist = 0
+    mutableState.update { state ->
+      val threadIndex = findThreadIndex(
+        projects = state.projects,
+        normalizedPath = normalizedPath,
+        provider = provider,
+        threadId = threadId,
+      ) ?: return@update state
+      val currentVisible = state.visibleThreadCounts[normalizedPath] ?: treeUiState.getVisibleThreadCount(normalizedPath)
+      if (threadIndex < currentVisible) {
+        return@update state
+      }
+      val minVisible = threadIndex + 1
+      var nextVisible = currentVisible
+      while (nextVisible < minVisible) {
+        nextVisible += DEFAULT_VISIBLE_THREAD_COUNT
+      }
+      deltaToPersist = nextVisible - currentVisible
+      state.copy(visibleThreadCounts = state.visibleThreadCounts + (normalizedPath to nextVisible))
+    }
+    if (deltaToPersist > 0) {
+      treeUiState.incrementVisibleThreadCount(normalizedPath, deltaToPersist)
+    }
+  }
+
+  fun buildInitialVisibleThreadCounts(knownPaths: List<String>): Map<String, Int> {
+    return buildInitialVisibleThreadCounts(
+      knownPaths = knownPaths,
+      currentVisibleThreadCounts = mutableState.value.visibleThreadCounts,
+    )
+  }
+
+  fun updateProject(path: String, update: (AgentProjectSessions) -> AgentProjectSessions) {
+    mutableState.update { state ->
+      val next = state.projects.map { project ->
+        if (project.path == path) update(project) else project
+      }
+      state.copy(projects = next, lastUpdatedAt = System.currentTimeMillis())
+    }
+  }
+
+  fun updateWorktree(projectPath: String, worktreePath: String, update: (AgentWorktree) -> AgentWorktree) {
+    mutableState.update { state ->
+      val next = state.projects.map { project ->
+        if (project.path == projectPath) {
+          project.copy(worktrees = project.worktrees.map { wt ->
+            if (wt.path == worktreePath) update(wt) else wt
+          })
+        }
+        else {
+          project
+        }
+      }
+      state.copy(projects = next, lastUpdatedAt = System.currentTimeMillis())
+    }
+  }
+
+  fun findWorktreeBranch(path: String): String? {
+    for (project in mutableState.value.projects) {
+      for (worktree in project.worktrees) {
+        if (worktree.path == path) return worktree.branch
+      }
+    }
+    return null
+  }
+
+  private fun buildInitialVisibleThreadCounts(
+    knownPaths: List<String>,
+    currentVisibleThreadCounts: Map<String, Int>,
+  ): Map<String, Int> {
+    val normalizedKnownPaths = knownPaths.mapTo(LinkedHashSet()) { normalizePath(it) }
+    val visibleThreadCounts = LinkedHashMap<String, Int>()
+    currentVisibleThreadCounts.forEach { (path, count) ->
+      val normalized = normalizePath(path)
+      if (normalized in normalizedKnownPaths && count > DEFAULT_VISIBLE_THREAD_COUNT) {
+        visibleThreadCounts[normalized] = count
+      }
+    }
+    for (path in normalizedKnownPaths) {
+      if (path in visibleThreadCounts) continue
+      val persisted = treeUiState.getVisibleThreadCount(path)
+      if (persisted > DEFAULT_VISIBLE_THREAD_COUNT) {
+        visibleThreadCounts[path] = persisted
+      }
+    }
+    return visibleThreadCounts
+  }
+
+  private fun findThreadIndex(
+    projects: List<AgentProjectSessions>,
+    normalizedPath: String,
+    provider: AgentSessionProvider,
+    threadId: String,
+  ): Int? {
+    val projectThreads = projects.firstOrNull { it.path == normalizedPath }?.threads
+    if (projectThreads != null) {
+      val index = projectThreads.indexOfFirst { it.provider == provider && it.id == threadId }
+      if (index >= 0) return index
+    }
+
+    projects.forEach { project ->
+      val worktreeThreads = project.worktrees.firstOrNull { it.path == normalizedPath }?.threads ?: return@forEach
+      val index = worktreeThreads.indexOfFirst { it.provider == provider && it.id == threadId }
+      if (index >= 0) return index
+    }
+
+    return null
+  }
+
+  private fun normalizePath(path: String): String {
+    return try {
+      Path.of(path).invariantSeparatorsPathString
+    }
+    catch (_: InvalidPathException) {
+      path
+    }
+  }
+}
@@ -45,7 +45,7 @@ internal fun thread(id: String, updatedAt: Long, provider: AgentSessionProvider)

 internal suspend fun withService(
  sessionSourcesProvider: () -> List<AgentSessionSource>,
-  projectEntriesProvider: suspend () -> List<AgentSessionsService.ProjectEntry>,
+  projectEntriesProvider: suspend () -> List<ProjectEntry>,
  treeUiState: SessionsTreeUiState = InMemorySessionsTreeUiState(),
  action: suspend (AgentSessionsService) -> Unit,
 ) {
@@ -69,9 +69,9 @@ internal suspend fun withService(
 internal fun openProjectEntry(
  path: String,
  name: String,
-  worktrees: List<AgentSessionsService.WorktreeEntry> = emptyList(),
-): AgentSessionsService.ProjectEntry {
-  return AgentSessionsService.ProjectEntry(
+  worktrees: List<WorktreeEntry> = emptyList(),
+): ProjectEntry {
+  return ProjectEntry(
    path = path,
    name = name,
    project = openProjectProxy(name),
@@ -82,9 +82,9 @@ internal fun openProjectEntry(
 internal fun closedProjectEntry(
  path: String,
  name: String,
-  worktrees: List<AgentSessionsService.WorktreeEntry> = emptyList(),
-): AgentSessionsService.ProjectEntry {
-  return AgentSessionsService.ProjectEntry(
+  worktrees: List<WorktreeEntry> = emptyList(),
+): ProjectEntry {
+  return ProjectEntry(
    path = path,
    name = name,
    project = null,
@@ -216,7 +216,7 @@ class AgentSessionsServiceOnDemandIntegrationTest {
            PROJECT_PATH,
            "Project A",
            worktrees = listOf(
-              AgentSessionsService.WorktreeEntry(
+              WorktreeEntry(
                path = WORKTREE_PATH,
                name = "project-feature",
                branch = "feature",