diff --git a/plugins/agent-workbench/claude/sessions/resources/intellij.agent.workbench.claude.sessions.xml b/plugins/agent-workbench/claude/sessions/resources/intellij.agent.workbench.claude.sessions.xml index b014541cb4ff..6c11091ec579 100644 --- a/plugins/agent-workbench/claude/sessions/resources/intellij.agent.workbench.claude.sessions.xml +++ b/plugins/agent-workbench/claude/sessions/resources/intellij.agent.workbench.claude.sessions.xml @@ -7,6 +7,8 @@ - + diff --git a/plugins/agent-workbench/claude/sessions/src/ClaudeAgentSessionProviderBridge.kt b/plugins/agent-workbench/claude/sessions/src/ClaudeAgentSessionProviderBridge.kt index 182eb190d9d7..31d42086d9d0 100644 --- a/plugins/agent-workbench/claude/sessions/src/ClaudeAgentSessionProviderBridge.kt +++ b/plugins/agent-workbench/claude/sessions/src/ClaudeAgentSessionProviderBridge.kt @@ -11,9 +11,6 @@ import com.intellij.agent.workbench.sessions.providers.AgentSessionSource internal class ClaudeAgentSessionProviderBridge( override val sessionSource: AgentSessionSource = ClaudeSessionSource(), ) : AgentSessionProviderBridge { - override val order: Int - get() = 10 - override val provider: AgentSessionProvider get() = AgentSessionProvider.CLAUDE diff --git a/plugins/agent-workbench/codex/sessions/resources/intellij.agent.workbench.codex.sessions.xml b/plugins/agent-workbench/codex/sessions/resources/intellij.agent.workbench.codex.sessions.xml index 6bd8b4c35969..7c404dd4e6b1 100644 --- a/plugins/agent-workbench/codex/sessions/resources/intellij.agent.workbench.codex.sessions.xml +++ b/plugins/agent-workbench/codex/sessions/resources/intellij.agent.workbench.codex.sessions.xml @@ -8,7 +8,9 @@ - + - diff --git a/plugins/agent-workbench/codex/sessions/src/CodexAgentSessionProviderBridge.kt b/plugins/agent-workbench/codex/sessions/src/CodexAgentSessionProviderBridge.kt index cdf238ca424d..b72e7568c9be 100644 --- a/plugins/agent-workbench/codex/sessions/src/CodexAgentSessionProviderBridge.kt +++ b/plugins/agent-workbench/codex/sessions/src/CodexAgentSessionProviderBridge.kt @@ -15,9 +15,6 @@ import com.intellij.openapi.components.service internal class CodexAgentSessionProviderBridge( override val sessionSource: AgentSessionSource = CodexSessionSource(), ) : AgentSessionProviderBridge { - override val order: Int - get() = 20 - override val provider: AgentSessionProvider get() = AgentSessionProvider.CODEX diff --git a/plugins/agent-workbench/sessions/src/AgentSessionModels.kt b/plugins/agent-workbench/sessions/src/AgentSessionModels.kt index 49a3d31e82aa..f7b4f7957dfe 100644 --- a/plugins/agent-workbench/sessions/src/AgentSessionModels.kt +++ b/plugins/agent-workbench/sessions/src/AgentSessionModels.kt @@ -50,6 +50,7 @@ data class AgentSessionThreadPreview( @JvmField val id: String, @JvmField val title: String, @JvmField val updatedAt: Long, + val provider: AgentSessionProvider = AgentSessionProvider.CODEX, ) @Immutable diff --git a/plugins/agent-workbench/sessions/src/AgentSessionsService.kt b/plugins/agent-workbench/sessions/src/AgentSessionsService.kt index a88e3df1c5ca..6d631cc1a8a3 100644 --- a/plugins/agent-workbench/sessions/src/AgentSessionsService.kt +++ b/plugins/agent-workbench/sessions/src/AgentSessionsService.kt @@ -66,6 +66,7 @@ internal class AgentSessionsService private constructor( private val serviceScope: CoroutineScope, private val sessionSourcesProvider: () -> List, private val projectEntriesProvider: suspend (AgentSessionsService) -> List, + private val treeUiState: SessionsTreeUiState, subscribeToProjectLifecycle: Boolean, ) { @Suppress("unused") @@ -73,6 +74,7 @@ internal class AgentSessionsService private constructor( serviceScope = serviceScope, sessionSourcesProvider = AgentSessionProviderBridges::sessionSources, projectEntriesProvider = { service -> service.collectProjects() }, + treeUiState = service(), subscribeToProjectLifecycle = true, ) @@ -80,11 +82,13 @@ internal class AgentSessionsService private constructor( serviceScope: CoroutineScope, sessionSourcesProvider: () -> List, projectEntriesProvider: suspend () -> List, + treeUiState: SessionsTreeUiState = InMemorySessionsTreeUiState(), subscribeToProjectLifecycle: Boolean = false, ) : this( serviceScope = serviceScope, sessionSourcesProvider = sessionSourcesProvider, projectEntriesProvider = { _ -> projectEntriesProvider() }, + treeUiState = treeUiState, subscribeToProjectLifecycle = subscribeToProjectLifecycle, ) @@ -280,49 +284,82 @@ internal class AgentSessionsService private constructor( } try { val entries = projectEntriesProvider(this@AgentSessionsService) - val currentProjectsByPath = mutableState.value.projects.associateBy { it.path } + val currentState = mutableState.value + 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 = buildInitialVisibleThreadCounts( + knownPaths = knownPaths, + currentVisibleThreadCounts = currentState.visibleThreadCounts, + ) val initialProjects = entries.map { entry -> - val existing = currentProjectsByPath[entry.path] + 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 = entry.path, + path = normalizedEntryPath, name = entry.name, branch = entry.branch ?: existing?.branch, isOpen = entry.project != null, isLoading = entry.project != null, - hasLoaded = existing?.hasLoaded ?: false, + hasLoaded = existing?.hasLoaded ?: (cachedPreviews != null), hasUnknownThreadCount = existing?.hasUnknownThreadCount ?: false, - threads = existing?.threads ?: emptyList(), + threads = existing?.threads ?: cachedThreads, errorMessage = existing?.errorMessage, providerWarnings = existing?.providerWarnings ?: emptyList(), worktrees = entry.worktreeEntries.map { wt -> - val existingWt = existing?.worktrees?.firstOrNull { it.path == wt.path } + 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 = wt.path, + path = normalizedWorktreePath, name = wt.name, branch = wt.branch, isOpen = wt.project != null, - isLoading = wt.project != null && hasExistingData, - hasLoaded = existingWt?.hasLoaded ?: false, + isLoading = wt.project != null && (hasExistingData || cachedWorktreePreviews != null), + hasLoaded = existingWt?.hasLoaded ?: (cachedWorktreePreviews != null), hasUnknownThreadCount = existingWt?.hasUnknownThreadCount ?: false, - threads = existingWt?.threads ?: emptyList(), + threads = existingWt?.threads ?: cachedWorktreeThreads, errorMessage = existingWt?.errorMessage, providerWarnings = existingWt?.providerWarnings ?: emptyList(), ) }, ) } - mutableState.update { it.copy(projects = initialProjects, lastUpdatedAt = System.currentTimeMillis()) } - - // Collect all open paths for prefetching (batch Codex calls) - val sessionSources = sessionSourcesProvider() - val openPaths = entries.flatMap { entry -> - buildList { - if (entry.project != null) add(entry.path) - entry.worktreeEntries.filter { it.project != null }.forEach { add(it.path) } - } + mutableState.update { + it.copy( + projects = initialProjects, + visibleThreadCounts = initialVisibleThreadCounts, + lastUpdatedAt = System.currentTimeMillis(), + ) } + val sessionSources = sessionSourcesProvider() + // Prefetch from all sources in parallel val prefetchedByProvider = coroutineScope { sessionSources.map { source -> @@ -342,8 +379,9 @@ internal class AgentSessionsService private constructor( coroutineScope { for (entry in entries) { launch { + val normalizedEntryPath = normalizePath(entry.path) if (entry.project == null) { - updateProject(entry.path) { it.copy(isLoading = false) } + updateProject(normalizedEntryPath) { it.copy(isLoading = false) } return@launch } val sourceResults = java.util.concurrent.CopyOnWriteArrayList() @@ -351,9 +389,9 @@ internal class AgentSessionsService private constructor( for (source in sessionSources) { launch { val sourceResult = try { - val prefetched = prefetchedByProvider[source.provider]?.get(entry.path) + val prefetched = prefetchedByProvider[source.provider]?.get(normalizedEntryPath) val threads = prefetched - ?: source.listThreadsFromOpenProject(path = entry.path, project = entry.project) + ?: source.listThreadsFromOpenProject(path = normalizedEntryPath, project = entry.project) AgentSessionSourceLoadResult( provider = source.provider, result = Result.success(threads), @@ -376,7 +414,7 @@ internal class AgentSessionsService private constructor( resolveWarningMessage = ::resolveProviderWarningMessage, ) val anySuccess = sourceResults.any { it.result.isSuccess } - updateProject(entry.path) { project -> + updateProject(normalizedEntryPath) { project -> project.copy( threads = partial.threads, providerWarnings = partial.providerWarnings, @@ -392,7 +430,7 @@ internal class AgentSessionsService private constructor( resolveErrorMessage = ::resolveErrorMessage, resolveWarningMessage = ::resolveProviderWarningMessage, ) - updateProject(entry.path) { project -> + updateProject(normalizedEntryPath) { project -> project.copy(isLoading = false, hasLoaded = true, hasUnknownThreadCount = finalResult.hasUnknownThreadCount, @@ -400,11 +438,16 @@ internal class AgentSessionsService private constructor( 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) { - updateWorktree(entry.path, wt.path) { it.copy(isLoading = false) } + updateWorktree(normalizedEntryPath, normalizedWorktreePath) { it.copy(isLoading = false) } return@launch } val sourceResults = java.util.concurrent.CopyOnWriteArrayList() @@ -412,9 +455,9 @@ internal class AgentSessionsService private constructor( for (source in sessionSources) { launch { val sourceResult = try { - val prefetched = prefetchedByProvider[source.provider]?.get(wt.path) + val prefetched = prefetchedByProvider[source.provider]?.get(normalizedWorktreePath) val threads = prefetched - ?: source.listThreadsFromOpenProject(path = wt.path, project = wt.project) + ?: source.listThreadsFromOpenProject(path = normalizedWorktreePath, project = wt.project) AgentSessionSourceLoadResult( provider = source.provider, result = Result.success(threads), @@ -436,7 +479,7 @@ internal class AgentSessionsService private constructor( resolveWarningMessage = ::resolveProviderWarningMessage, ) val anySuccess = sourceResults.any { it.result.isSuccess } - updateWorktree(entry.path, wt.path) { worktree -> + updateWorktree(normalizedEntryPath, normalizedWorktreePath) { worktree -> worktree.copy( threads = partial.threads, providerWarnings = partial.providerWarnings, @@ -451,7 +494,7 @@ internal class AgentSessionsService private constructor( resolveErrorMessage = ::resolveErrorMessage, resolveWarningMessage = ::resolveProviderWarningMessage, ) - updateWorktree(entry.path, wt.path) { worktree -> + updateWorktree(normalizedEntryPath, normalizedWorktreePath) { worktree -> worktree.copy(isLoading = false, hasLoaded = true, hasUnknownThreadCount = finalResult.hasUnknownThreadCount, @@ -459,6 +502,9 @@ internal class AgentSessionsService private constructor( errorMessage = finalResult.errorMessage, providerWarnings = finalResult.providerWarnings) } + if (finalResult.errorMessage == null) { + treeUiState.setOpenProjectThreadPreviews(normalizedWorktreePath, finalResult.threads.toThreadPreviews()) + } } } } @@ -510,14 +556,22 @@ internal class AgentSessionsService private constructor( } fun showMoreThreads(path: String) { + val normalizedPath = normalizePath(path) + var deltaToPersist = 0 mutableState.update { state -> - val current = state.visibleThreadCounts[path] ?: DEFAULT_VISIBLE_THREAD_COUNT - state.copy(visibleThreadCounts = state.visibleThreadCounts + (path to (current + DEFAULT_VISIBLE_THREAD_COUNT))) + 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, @@ -525,7 +579,7 @@ internal class AgentSessionsService private constructor( provider = provider, threadId = threadId, ) ?: return@update state - val currentVisible = state.visibleThreadCounts[normalizedPath] ?: DEFAULT_VISIBLE_THREAD_COUNT + val currentVisible = state.visibleThreadCounts[normalizedPath] ?: treeUiState.getVisibleThreadCount(normalizedPath) if (threadIndex < currentVisible) { return@update state } @@ -534,8 +588,12 @@ internal class AgentSessionsService private constructor( 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 openChatThread(path: String, thread: AgentSessionThread, currentProject: Project? = null) { @@ -1396,6 +1454,51 @@ internal class AgentSessionsService private constructor( return null } + private fun buildInitialVisibleThreadCounts( + knownPaths: List, + currentVisibleThreadCounts: Map, + ): Map { + val normalizedKnownPaths = knownPaths.mapTo(LinkedHashSet()) { normalizePath(it) } + val visibleThreadCounts = LinkedHashMap() + 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 List.toCachedSessionThreads(): List { + return map { preview -> + AgentSessionThread( + id = preview.id, + title = preview.title, + updatedAt = preview.updatedAt, + archived = false, + provider = preview.provider, + ) + } + } + + private fun List.toThreadPreviews(): List { + return map { thread -> + AgentSessionThreadPreview( + id = thread.id, + title = thread.title, + updatedAt = thread.updatedAt, + provider = thread.provider, + ) + } + } + private data class ProviderRefreshOutcome( val threads: List? = null, val warningMessage: String? = null, diff --git a/plugins/agent-workbench/sessions/src/AgentSessionsToolWindow.kt b/plugins/agent-workbench/sessions/src/AgentSessionsToolWindow.kt index 9fbe2800d9a9..62202375ccc0 100644 --- a/plugins/agent-workbench/sessions/src/AgentSessionsToolWindow.kt +++ b/plugins/agent-workbench/sessions/src/AgentSessionsToolWindow.kt @@ -84,6 +84,7 @@ internal fun agentSessionsToolWindow(currentProject: Project) { onOpenThread = { path, thread -> service.openChatThread(path, thread, currentProject) }, onOpenSubAgent = { path, thread, subAgent -> service.openChatSubAgent(path, thread, subAgent, currentProject) }, onCreateSession = { path, provider, mode -> service.createNewSession(path, provider, mode, currentProject) }, + treeUiState = uiStateService, lastUsedProvider = lastUsedProvider, visibleProjectCount = state.visibleProjectCount, onShowMoreProjects = { service.showMoreProjects() }, @@ -111,6 +112,7 @@ internal fun agentSessionsToolWindowContent( onOpenThread: (String, AgentSessionThread) -> Unit = { _, _ -> }, onOpenSubAgent: (String, AgentSessionThread, AgentSubAgent) -> Unit = { _, _, _ -> }, onCreateSession: (String, AgentSessionProvider, AgentSessionLaunchMode) -> Unit = { _, _, _ -> }, + treeUiState: SessionsTreeUiState? = null, lastUsedProvider: AgentSessionProvider? = null, nowProvider: () -> Long = { System.currentTimeMillis() }, visibleProjectCount: Int = Int.MAX_VALUE, @@ -122,6 +124,7 @@ internal fun agentSessionsToolWindowContent( onEnableClaudeQuotaWidget: () -> Unit = {}, onDismissClaudeQuotaHint: () -> Unit = {}, ) { + val effectiveTreeUiState = treeUiState ?: remember { InMemorySessionsTreeUiState() } Column( modifier = Modifier .fillMaxSize() @@ -146,6 +149,7 @@ internal fun agentSessionsToolWindowContent( onOpenThread = onOpenThread, onOpenSubAgent = onOpenSubAgent, onCreateSession = onCreateSession, + treeUiState = effectiveTreeUiState, lastUsedProvider = lastUsedProvider, nowProvider = nowProvider, visibleProjectCount = visibleProjectCount, diff --git a/plugins/agent-workbench/sessions/src/AgentSessionsTreeUiStateService.kt b/plugins/agent-workbench/sessions/src/AgentSessionsTreeUiStateService.kt index a79b6481db8d..ab7f21445e1a 100644 --- a/plugins/agent-workbench/sessions/src/AgentSessionsTreeUiStateService.kt +++ b/plugins/agent-workbench/sessions/src/AgentSessionsTreeUiStateService.kt @@ -158,6 +158,9 @@ internal class AgentSessionsTreeUiStateService id = preview.id, title = preview.title, updatedAt = preview.updatedAt, + provider = preview.provider + ?.let(AgentSessionProvider::fromOrNull) + ?: AgentSessionProvider.CODEX, ) } } @@ -165,7 +168,14 @@ internal class AgentSessionsTreeUiStateService override fun setOpenProjectThreadPreviews(path: String, threads: List): Boolean { val normalizedPath = normalizeSessionsProjectPath(path) val normalizedPreviews = normalizeOpenProjectThreadPreviewList(threads) - .map { thread -> ThreadPreviewState(id = thread.id, title = thread.title, updatedAt = thread.updatedAt) } + .map { thread -> + ThreadPreviewState( + id = thread.id, + title = thread.title, + updatedAt = thread.updatedAt, + provider = thread.provider.value, + ) + } val current = state.openProjectThreadPreviewsByProject[normalizedPath] if (current == normalizedPreviews) { return false @@ -278,6 +288,7 @@ internal class AgentSessionsTreeUiStateService @JvmField val id: String, @JvmField val title: String, @JvmField val updatedAt: Long, + @JvmField val provider: String? = null, ) } diff --git a/plugins/agent-workbench/sessions/src/SessionTree.kt b/plugins/agent-workbench/sessions/src/SessionTree.kt index 11155336b719..c7f5d0f4cb18 100644 --- a/plugins/agent-workbench/sessions/src/SessionTree.kt +++ b/plugins/agent-workbench/sessions/src/SessionTree.kt @@ -32,6 +32,7 @@ internal fun sessionTree( onOpenThread: (String, AgentSessionThread) -> Unit, onOpenSubAgent: (String, AgentSessionThread, AgentSubAgent) -> Unit, onCreateSession: (String, AgentSessionProvider, AgentSessionLaunchMode) -> Unit = { _, _, _ -> }, + treeUiState: SessionsTreeUiState, lastUsedProvider: AgentSessionProvider? = null, nowProvider: () -> Long, visibleProjectCount: Int = Int.MAX_VALUE, @@ -41,12 +42,17 @@ internal fun sessionTree( selectedTreeId: SessionTreeId? = null, ) { val stateHolder = rememberSessionTreeStateHolder( - onProjectExpanded = onProjectExpanded, - onProjectCollapsed = {}, + onProjectExpanded = { path -> + treeUiState.setProjectCollapsed(path, collapsed = false) + onProjectExpanded(path) + }, + onProjectCollapsed = { path -> + treeUiState.setProjectCollapsed(path, collapsed = true) + }, onWorktreeExpanded = onWorktreeExpanded, ) val treeState = stateHolder.treeState - val autoOpenNodes = remember(projects, visibleProjectCount) { + val autoOpenNodes = remember(projects, visibleProjectCount, treeUiState) { projects.take(visibleProjectCount) .filter { it.isOpen || @@ -54,6 +60,7 @@ internal fun sessionTree( it.providerWarnings.isNotEmpty() || it.worktrees.any { wt -> wt.isOpen } } + .filterNot { treeUiState.isProjectCollapsed(it.path) } .map { SessionTreeId.Project(it.path) } } LaunchedEffect(autoOpenNodes) { diff --git a/plugins/agent-workbench/sessions/src/providers/AgentSessionProviderBridge.kt b/plugins/agent-workbench/sessions/src/providers/AgentSessionProviderBridge.kt index f90e188e1396..9cbb80720180 100644 --- a/plugins/agent-workbench/sessions/src/providers/AgentSessionProviderBridge.kt +++ b/plugins/agent-workbench/sessions/src/providers/AgentSessionProviderBridge.kt @@ -5,9 +5,6 @@ import com.intellij.agent.workbench.sessions.AgentSessionLaunchMode import com.intellij.agent.workbench.sessions.AgentSessionProvider interface AgentSessionProviderBridge { - val order: Int - get() = 0 - val provider: AgentSessionProvider val displayNameKey: String val newSessionLabelKey: String diff --git a/plugins/agent-workbench/sessions/src/providers/AgentSessionSourcesRegistry.kt b/plugins/agent-workbench/sessions/src/providers/AgentSessionSourcesRegistry.kt index 5f6bdf70fba1..f03bd1644351 100644 --- a/plugins/agent-workbench/sessions/src/providers/AgentSessionSourcesRegistry.kt +++ b/plugins/agent-workbench/sessions/src/providers/AgentSessionSourcesRegistry.kt @@ -50,11 +50,10 @@ private fun buildAgentSessionProviderBridgeSnapshot(bridges: Iterable List, projectEntriesProvider: suspend () -> List, + treeUiState: SessionsTreeUiState = InMemorySessionsTreeUiState(), action: suspend (AgentSessionsService) -> Unit, ) { @Suppress("RAW_SCOPE_CREATION") @@ -55,6 +56,7 @@ internal suspend fun withService( serviceScope = scope, sessionSourcesProvider = sessionSourcesProvider, projectEntriesProvider = projectEntriesProvider, + treeUiState = treeUiState, subscribeToProjectLifecycle = false, ) action(service) diff --git a/plugins/agent-workbench/sessions/testSrc/AgentSessionsServiceOnDemandIntegrationTest.kt b/plugins/agent-workbench/sessions/testSrc/AgentSessionsServiceOnDemandIntegrationTest.kt index 08fd9246f531..1d540db48ab2 100644 --- a/plugins/agent-workbench/sessions/testSrc/AgentSessionsServiceOnDemandIntegrationTest.kt +++ b/plugins/agent-workbench/sessions/testSrc/AgentSessionsServiceOnDemandIntegrationTest.kt @@ -12,6 +12,7 @@ import java.util.concurrent.atomic.AtomicInteger class AgentSessionsServiceOnDemandIntegrationTest { @Test fun ensureThreadVisibleExpandsProjectVisibleCountForHiddenThread() = runBlocking { + val treeUiState = InMemorySessionsTreeUiState() withService( sessionSourcesProvider = { listOf( @@ -38,6 +39,7 @@ class AgentSessionsServiceOnDemandIntegrationTest { projectEntriesProvider = { listOf(closedProjectEntry(PROJECT_PATH, "Project A")) }, + treeUiState = treeUiState, ) { service -> service.refresh() waitForCondition { @@ -53,6 +55,39 @@ class AgentSessionsServiceOnDemandIntegrationTest { assertThat(service.state.value.visibleThreadCounts[PROJECT_PATH]) .isEqualTo(DEFAULT_VISIBLE_THREAD_COUNT + DEFAULT_VISIBLE_THREAD_COUNT) + assertThat(treeUiState.getVisibleThreadCount(PROJECT_PATH)) + .isEqualTo(DEFAULT_VISIBLE_THREAD_COUNT + DEFAULT_VISIBLE_THREAD_COUNT) + } + } + + @Test + fun showMoreThreadsPersistsVisibleCountToUiState() = runBlocking { + val treeUiState = InMemorySessionsTreeUiState() + withService( + sessionSourcesProvider = { + listOf( + ScriptedSessionSource( + provider = AgentSessionProvider.CODEX, + canReportExactThreadCount = true, + ), + ) + }, + projectEntriesProvider = { + listOf(closedProjectEntry(PROJECT_PATH, "Project A")) + }, + treeUiState = treeUiState, + ) { service -> + service.refresh() + waitForCondition { + service.state.value.projects.any { it.path == PROJECT_PATH } + } + + service.showMoreThreads(PROJECT_PATH) + + assertThat(service.state.value.visibleThreadCounts[PROJECT_PATH]) + .isEqualTo(DEFAULT_VISIBLE_THREAD_COUNT + DEFAULT_VISIBLE_THREAD_COUNT) + assertThat(treeUiState.getVisibleThreadCount(PROJECT_PATH)) + .isEqualTo(DEFAULT_VISIBLE_THREAD_COUNT + DEFAULT_VISIBLE_THREAD_COUNT) } } diff --git a/plugins/agent-workbench/sessions/testSrc/AgentSessionsServiceRefreshIntegrationTest.kt b/plugins/agent-workbench/sessions/testSrc/AgentSessionsServiceRefreshIntegrationTest.kt index 8c8a60dfe6f0..ed1bf5c7abf1 100644 --- a/plugins/agent-workbench/sessions/testSrc/AgentSessionsServiceRefreshIntegrationTest.kt +++ b/plugins/agent-workbench/sessions/testSrc/AgentSessionsServiceRefreshIntegrationTest.kt @@ -2,6 +2,7 @@ package com.intellij.agent.workbench.sessions import com.intellij.testFramework.junit5.TestApplication +import kotlinx.coroutines.CompletableDeferred import kotlinx.coroutines.flow.MutableSharedFlow import kotlinx.coroutines.runBlocking import org.assertj.core.api.Assertions.assertThat @@ -9,6 +10,92 @@ import org.junit.jupiter.api.Test @TestApplication class AgentSessionsServiceRefreshIntegrationTest { + @Test + fun refreshShowsCachedOpenProjectThreadsBeforeProviderLoadCompletes() = runBlocking { + val started = CompletableDeferred() + val release = CompletableDeferred() + val treeUiState = InMemorySessionsTreeUiState() + treeUiState.setOpenProjectThreadPreviews( + PROJECT_PATH, + listOf( + AgentSessionThreadPreview( + id = "cached-1", + title = "Cached", + updatedAt = 100, + provider = AgentSessionProvider.CLAUDE, + ) + ), + ) + + withService( + sessionSourcesProvider = { + listOf( + ScriptedSessionSource( + provider = AgentSessionProvider.CLAUDE, + listFromOpenProject = { path, _ -> + if (path != PROJECT_PATH) { + emptyList() + } + else { + started.complete(Unit) + release.await() + listOf(thread(id = "claude-1", updatedAt = 200, provider = AgentSessionProvider.CLAUDE)) + } + }, + ) + ) + }, + projectEntriesProvider = { + listOf(openProjectEntry(PROJECT_PATH, "Project A")) + }, + treeUiState = treeUiState, + ) { service -> + service.refresh() + started.await() + + waitForCondition { + val project = service.state.value.projects.firstOrNull { it.path == PROJECT_PATH } ?: return@waitForCondition false + project.hasLoaded && project.threads.map { it.id } == listOf("cached-1") + } + + release.complete(Unit) + waitForCondition { + service.state.value.projects.firstOrNull { it.path == PROJECT_PATH }?.threads?.map { it.id } == listOf("claude-1") + } + + assertThat(treeUiState.getOpenProjectThreadPreviews(PROJECT_PATH).orEmpty().map { it.id }) + .containsExactly("claude-1") + } + } + + @Test + fun refreshRestoresPersistedVisibleThreadCountForKnownPath() = runBlocking { + val treeUiState = InMemorySessionsTreeUiState() + treeUiState.incrementVisibleThreadCount(PROJECT_PATH, delta = 6) + + withService( + sessionSourcesProvider = { + listOf( + ScriptedSessionSource( + provider = AgentSessionProvider.CODEX, + ) + ) + }, + projectEntriesProvider = { + listOf(closedProjectEntry(PROJECT_PATH, "Project A")) + }, + treeUiState = treeUiState, + ) { service -> + service.refresh() + waitForCondition { + service.state.value.projects.any { it.path == PROJECT_PATH } + } + + assertThat(service.state.value.visibleThreadCounts[PROJECT_PATH]) + .isEqualTo(DEFAULT_VISIBLE_THREAD_COUNT + 6) + } + } + @Test fun refreshMergesMixedProviderThreadsAndMarksUnknownCount() = runBlocking { withService( diff --git a/plugins/agent-workbench/sessions/testSrc/AgentSessionsToolWindowTest.kt b/plugins/agent-workbench/sessions/testSrc/AgentSessionsToolWindowTest.kt index 1e23924760aa..92bb97f590bd 100644 --- a/plugins/agent-workbench/sessions/testSrc/AgentSessionsToolWindowTest.kt +++ b/plugins/agent-workbench/sessions/testSrc/AgentSessionsToolWindowTest.kt @@ -398,6 +398,99 @@ class AgentSessionsToolWindowTest { .assertIsDisplayed() } + @Test + fun collapsedProjectInUiStateIsNotAutoExpandedAndExpandingClearsCollapsedFlag() { + val projectPath = "/work/project-a" + val treeUiState = InMemorySessionsTreeUiState() + treeUiState.setProjectCollapsed(projectPath, collapsed = true) + val projects = listOf( + AgentProjectSessions( + path = projectPath, + name = "Project A", + isOpen = true, + hasLoaded = true, + ), + ) + + composeRule.setContentWithTheme { + agentSessionsToolWindowContent( + state = AgentSessionsState(projects = projects), + onRefresh = {}, + onOpenProject = {}, + treeUiState = treeUiState, + ) + } + + composeRule.onNodeWithText("Project A").assertIsDisplayed() + composeRule.onAllNodesWithText(AgentSessionsBundle.message("toolwindow.empty.project")) + .assertCountEquals(0) + + composeRule.onNodeWithText("Project A") + .assertIsDisplayed() + .performMouseInput { doubleClick() } + + composeRule.onNodeWithText(AgentSessionsBundle.message("toolwindow.empty.project")) + .assertIsDisplayed() + composeRule.runOnIdle { + assertThat(treeUiState.isProjectCollapsed(projectPath)).isFalse() + } + } + + @Test + fun collapsedProjectInPersistentUiStateRemainsCollapsedAfterContentRefresh() { + val projectPath = "/work/project-a" + val treeUiState = AgentSessionsTreeUiStateService() + val projects = listOf( + AgentProjectSessions( + path = projectPath, + name = "Project A", + isOpen = true, + hasLoaded = true, + ), + ) + + composeRule.setContentWithTheme { + agentSessionsToolWindowContent( + state = AgentSessionsState(projects = projects, lastUpdatedAt = 1L), + onRefresh = {}, + onOpenProject = {}, + treeUiState = treeUiState, + ) + } + + composeRule.onNodeWithText(AgentSessionsBundle.message("toolwindow.empty.project")) + .assertIsDisplayed() + composeRule.onNodeWithText("Project A") + .assertIsDisplayed() + .performMouseInput { doubleClick() } + composeRule.onAllNodesWithText(AgentSessionsBundle.message("toolwindow.empty.project")) + .assertCountEquals(0) + composeRule.runOnIdle { + assertThat(treeUiState.isProjectCollapsed(projectPath)).isTrue() + } + + composeRule.setContentWithTheme { + agentSessionsToolWindowContent( + state = AgentSessionsState(projects = projects, lastUpdatedAt = 2L), + onRefresh = {}, + onOpenProject = {}, + treeUiState = treeUiState, + ) + } + + composeRule.onNodeWithText("Project A").assertIsDisplayed() + composeRule.onAllNodesWithText(AgentSessionsBundle.message("toolwindow.empty.project")) + .assertCountEquals(0) + composeRule.onNodeWithText("Project A") + .assertIsDisplayed() + .performMouseInput { doubleClick() } + composeRule.onNodeWithText(AgentSessionsBundle.message("toolwindow.empty.project")) + .assertIsDisplayed() + composeRule.runOnIdle { + assertThat(treeUiState.isProjectCollapsed(projectPath)).isFalse() + } + } + @Test fun expandingLoadedEmptyProjectShowsEmptyChildRow() { val now = 1_700_000_000_000L diff --git a/plugins/agent-workbench/sessions/testSrc/AgentSessionsTreeUiStateServiceTest.kt b/plugins/agent-workbench/sessions/testSrc/AgentSessionsTreeUiStateServiceTest.kt index 2557ce9fb694..9031db0e095a 100644 --- a/plugins/agent-workbench/sessions/testSrc/AgentSessionsTreeUiStateServiceTest.kt +++ b/plugins/agent-workbench/sessions/testSrc/AgentSessionsTreeUiStateServiceTest.kt @@ -61,6 +61,59 @@ class AgentSessionsTreeUiStateServiceTest { assertNull(uiState.getOpenProjectThreadPreviews("/work/project-a")) } + @Test + fun openProjectThreadPreviewCachePreservesProvider() { + val uiState = AgentSessionsTreeUiStateService() + val threads = listOf( + AgentSessionThreadPreview( + id = "claude-thread", + title = "Claude Thread", + updatedAt = 5L, + provider = AgentSessionProvider.CLAUDE, + ) + ) + + assertTrue(uiState.setOpenProjectThreadPreviews("/work/project-a", threads)) + assertEquals( + AgentSessionProvider.CLAUDE, + uiState.getOpenProjectThreadPreviews("/work/project-a")?.single()?.provider, + ) + } + + @Test + fun treeStateFieldsSurviveServiceStateRoundTrip() { + val original = AgentSessionsTreeUiStateService() + val path = "/work/project-a/" + original.setProjectCollapsed(path, collapsed = true) + original.incrementVisibleThreadCount(path, delta = 4) + original.setOpenProjectThreadPreviews( + path, + listOf( + AgentSessionThreadPreview( + id = "claude-thread", + title = "Claude Thread", + updatedAt = 20, + provider = AgentSessionProvider.CLAUDE, + ), + AgentSessionThreadPreview( + id = "codex-thread", + title = "Codex Thread", + updatedAt = 10, + provider = AgentSessionProvider.CODEX, + ), + ) + ) + + val reloaded = AgentSessionsTreeUiStateService() + reloaded.loadState(original.state) + + assertTrue(reloaded.isProjectCollapsed("/work/project-a")) + assertEquals(DEFAULT_VISIBLE_THREAD_COUNT + 4, reloaded.getVisibleThreadCount("/work/project-a")) + val previews = reloaded.getOpenProjectThreadPreviews("/work/project-a").orEmpty() + assertEquals(listOf("claude-thread", "codex-thread"), previews.map { it.id }) + assertEquals(listOf(AgentSessionProvider.CLAUDE, AgentSessionProvider.CODEX), previews.map { it.provider }) + } + @Test fun lastUsedProviderDefaultsToNull() { val uiState = AgentSessionsTreeUiStateService() diff --git a/plugins/agent-workbench/spec/agent-sessions-testing.spec.md b/plugins/agent-workbench/spec/agent-sessions-testing.spec.md index 215e65ca7834..5dd1ef71a0f7 100644 --- a/plugins/agent-workbench/spec/agent-sessions-testing.spec.md +++ b/plugins/agent-workbench/spec/agent-sessions-testing.spec.md @@ -2,6 +2,7 @@ name: Agent Threads Testing description: Coverage requirements for provider aggregation, tree rendering, and backend contracts in Agent Threads. targets: + - ../sessions/src/AgentSessionsTreeUiStateService.kt - ../codex/sessions/testSrc/CodexRolloutSessionBackendTest.kt - ../codex/sessions/testSrc/CodexSessionBackendSelectorTest.kt - ../sessions/testSrc/AgentSessionLoadAggregationTest.kt @@ -10,6 +11,7 @@ targets: - ../sessions/testSrc/AgentSessionsServiceConcurrencyIntegrationTest.kt - ../sessions/testSrc/AgentSessionsServiceIntegrationTestSupport.kt - ../sessions/testSrc/AgentSessionsToolWindowTest.kt + - ../sessions/testSrc/AgentSessionsTreeUiStateServiceTest.kt - ../sessions/testSrc/CodexSessionsPagingLogicTest.kt - ../sessions/testSrc/CodexAppServerClientTest.kt - ../sessions/testSrc/CodexAppServerClientTestSupport.kt @@ -19,7 +21,7 @@ targets: # Agent Threads Testing Status: Draft -Date: 2026-02-11 +Date: 2026-02-16 ## Summary Define required test coverage for the multi-provider Agent Threads stack: source aggregation, service behavior, tree/UI rendering, and Codex backend protocol compatibility. @@ -43,10 +45,14 @@ Define required test coverage for the multi-provider Agent Threads stack: source - Service integration tests must cover: - mixed-provider refresh merge, - provider warning and blocking error paths, - - unknown-count behavior when unknown provider fails/succeeds. + - unknown-count behavior when unknown provider fails/succeeds, + - cached preview rows rendered before open-path provider load completes, + - persisted visible thread count restoration during refresh bootstrap. - On-demand integration tests must cover: - project request deduplication, - - worktree request deduplication with refresh interaction. + - worktree request deduplication with refresh interaction, + - `showMoreThreads` visible-count persistence, + - `ensureThreadVisible` visible-count persistence. - Concurrency integration tests must verify refresh mutex deduplicates overlapping refresh calls. - Codex rollout backend tests must cover rollout parsing/activity behavior as the default thread-discovery path. - Codex backend selector tests must verify rollout default behavior and explicit app-server override behavior. @@ -54,7 +60,13 @@ Define required test coverage for the multi-provider Agent Threads stack: source - provider warning rendering, - error row precedence over warnings, - `More…` rendering for unknown count, - - `More (N)` rendering for exact count. + - `More (N)` rendering for exact count, + - persisted collapsed state blocking default auto-expand, + - collapsed-state persistence across content refresh/recreation when persistent tree UI state is used. +- Tree UI state service tests must cover: + - collapsed/visible-count/open-preview state round-trip, + - preview provider identity persistence, + - backward-compatible provider default for legacy preview entries with missing provider value. - Codex compatibility tests must cover cursor-loop/no-progress guard behavior in `seedInitialVisibleThreads`. - Codex app-server contract tests must run against mock backend always and real backend when available. @@ -63,6 +75,7 @@ Define required test coverage for the multi-provider Agent Threads stack: source [@test] ../sessions/testSrc/AgentSessionsServiceOnDemandIntegrationTest.kt [@test] ../sessions/testSrc/AgentSessionsServiceConcurrencyIntegrationTest.kt [@test] ../sessions/testSrc/AgentSessionsToolWindowTest.kt +[@test] ../sessions/testSrc/AgentSessionsTreeUiStateServiceTest.kt [@test] ../sessions/testSrc/CodexSessionsPagingLogicTest.kt [@test] ../sessions/testSrc/CodexAppServerClientTest.kt diff --git a/plugins/agent-workbench/spec/agent-sessions-thread-visibility.spec.md b/plugins/agent-workbench/spec/agent-sessions-thread-visibility.spec.md index a0e113cd5fa4..6a28bf5405aa 100644 --- a/plugins/agent-workbench/spec/agent-sessions-thread-visibility.spec.md +++ b/plugins/agent-workbench/spec/agent-sessions-thread-visibility.spec.md @@ -8,12 +8,13 @@ targets: - ../sessions/resources/messages/AgentSessionsBundle.properties - ../sessions/testSrc/AgentSessionsToolWindowTest.kt - ../sessions/testSrc/AgentSessionsServiceRefreshIntegrationTest.kt + - ../sessions/testSrc/AgentSessionsServiceOnDemandIntegrationTest.kt --- # Agent Threads Visibility and More Row Status: Draft -Date: 2026-02-11 +Date: 2026-02-16 ## Summary Define one consistent visibility model for project/worktree thread rows so empty state, warning/error rows, and `More` rows never conflict. @@ -29,6 +30,10 @@ Define one consistent visibility model for project/worktree thread rows so empty ## Requirements - Initial visible thread count per path is `DEFAULT_VISIBLE_THREAD_COUNT` (3). +- Visible thread count lookup order per normalized path is: + - in-memory runtime `visibleThreadCounts` entry, + - persisted tree UI state visible count, + - default value (`DEFAULT_VISIBLE_THREAD_COUNT`). - For project rows, show `More` row only when `project.threads.size > visibleCount`. - For worktree rows, show `More` row only when `worktree.threads.size > visibleCount`. - If `hasUnknownThreadCount=true` for the corresponding node, `More` row must render without count (`toolwindow.action.more`). @@ -41,15 +46,21 @@ Define one consistent visibility model for project/worktree thread rows so empty - no provider warnings. - Project/worktree error rows take precedence over provider warning and empty rows. - Clicking a `More` row must increase visible count for that path by +3. +- Clicking a `More` row must persist that +3 increment to tree UI state for the same normalized path. +- `ensureThreadVisible(path, provider, threadId)` must raise visible count in +3 steps until the target thread is visible and persist the increment. +- Refresh bootstrap must restore persisted visible counts greater than default for known project/worktree paths. +- Persisted visibility keys must use normalized paths so `/path` and `/path/` resolve to one entry. - Tree-side `More` click handling must not issue backend calls directly; it only updates visibility state. [@test] ../sessions/testSrc/AgentSessionsToolWindowTest.kt [@test] ../sessions/testSrc/AgentSessionsServiceRefreshIntegrationTest.kt +[@test] ../sessions/testSrc/AgentSessionsServiceOnDemandIntegrationTest.kt ## User Experience - Exact count case: `More (N)`. - Unknown count case: `More…`. - Empty row is muted helper text and mutually exclusive with `More` for the same node. +- Non-default per-path thread visibility must persist across refresh/reopen for the same normalized path. ## Data & Backend - Unknown-count state is produced by the service aggregation layer (`hasUnknownThreadCount`), not by tree rendering logic. diff --git a/plugins/agent-workbench/spec/agent-sessions.spec.md b/plugins/agent-workbench/spec/agent-sessions.spec.md index 5fe8f02405bc..c9808136ed41 100644 --- a/plugins/agent-workbench/spec/agent-sessions.spec.md +++ b/plugins/agent-workbench/spec/agent-sessions.spec.md @@ -14,7 +14,7 @@ targets: # Agent Threads Tool Window Status: Draft -Date: 2026-02-15 +Date: 2026-02-16 ## Summary Define the Agent Threads tool window as a provider-agnostic, project-scoped session browser. Threads from supported providers are aggregated per project/worktree, rendered in one tree, and opened through a shared chat routing flow. @@ -39,6 +39,18 @@ Define the Agent Threads tool window as a provider-agnostic, project-scoped sess - If at least one provider succeeds, successful threads must be shown and failed providers must surface provider-local warning rows. - If all providers fail for a project/worktree load, show blocking project/worktree error state and suppress provider warning rows for that load. - Unknown provider totals must propagate via `hasUnknownThreadCount` and drive unknown-count `More…` rendering. +- Sessions tree UI state must persist by normalized path: + - collapsed project/worktree state, + - per-path visible thread count, + - open-path thread preview cache. +- Refresh bootstrap must immediately seed open project/worktree nodes from cached previews when available and mark those paths loaded until live provider results arrive. +- Refresh bootstrap must restore persisted visible thread counts above default for known project/worktree paths. +- Refresh bootstrap must retain preview cache only for currently open project/worktree paths and prune stale closed-path entries. +- Final merged refresh results must update preview cache for a path only when that path does not end in blocking error. +- Auto-open default project expansion must skip paths persisted as collapsed. +- User collapse/expand interactions must update persisted collapsed state. +- Cached preview entries must preserve provider identity; missing legacy provider value must default to Codex for backward compatibility. +- `More` and programmatic visibility expansion (`ensureThreadVisible`) must persist visible-count increments in tree UI state. - On-demand project/worktree loading must deduplicate concurrent requests for the same path. - Concurrent refresh requests must be deduplicated while a refresh is already running. - Project primary click must open/focus the project; closed projects must expose `Open` in context menu. @@ -56,6 +68,7 @@ Define the Agent Threads tool window as a provider-agnostic, project-scoped sess [@test] ../sessions/testSrc/AgentSessionsServiceOnDemandIntegrationTest.kt [@test] ../sessions/testSrc/AgentSessionsServiceConcurrencyIntegrationTest.kt [@test] ../sessions/testSrc/AgentSessionsToolWindowTest.kt +[@test] ../sessions/testSrc/AgentSessionsTreeUiStateServiceTest.kt ## User Experience - Project rows are always expandable and may show worktree children.