diff --git a/plugins/agent-workbench/sessions-toolwindow/testSrc/AgentSessionsToolWindowFactorySwingTest.kt b/plugins/agent-workbench/sessions-toolwindow/testSrc/AgentSessionsToolWindowFactorySwingTest.kt
index b5213e8a1ccc..45b9359f88c1 100644
--- a/plugins/agent-workbench/sessions-toolwindow/testSrc/AgentSessionsToolWindowFactorySwingTest.kt
+++ b/plugins/agent-workbench/sessions-toolwindow/testSrc/AgentSessionsToolWindowFactorySwingTest.kt
@@ -406,6 +406,7 @@ class AgentSessionsToolWindowFactorySwingTest {
       .contains("AgentWorkbenchSessions.TreePopup.Open")
       .contains("AgentWorkbenchSessions.TreePopup.More")
       .contains(AgentWorkbenchActionIds.Sessions.TreePopup.NEW_THREAD)
+      .contains("AgentWorkbenchSessions.TreePopup.TaskFolderAgent")
       .contains("AgentWorkbenchSessions.TreePopup.Rename")
       .contains(AgentWorkbenchActionIds.Sessions.TreePopup.ARCHIVE)
       .contains("AgentWorkbenchSessions.TreePopup.Unarchive")
@@ -421,6 +422,7 @@ class AgentSessionsToolWindowFactorySwingTest {
     val copyThreadIdIndex = entries.requiredIndex("AgentWorkbenchSessions.TreePopup.CopyThreadId")
     val copyReferenceIndex = entries.requiredIndex("CopyReferencePopupGroup")
     val createTaskFolderIndex = entries.requiredIndex("AgentWorkbenchSessions.TreePopup.CreateTaskFolder")
+    val taskFolderAgentIndex = entries.requiredIndex("AgentWorkbenchSessions.TreePopup.TaskFolderAgent")
     val moveToTaskFolderIndex = entries.requiredIndex("AgentWorkbenchSessions.TreePopup.MoveToTaskFolder")
     val removeFromTaskFolderIndex = entries.requiredIndex("AgentWorkbenchSessions.TreePopup.RemoveFromTaskFolder")
     val markTaskFolderDoneIndex = entries.requiredIndex("AgentWorkbenchSessions.TreePopup.MarkTaskFolderDone")
@@ -430,7 +432,8 @@ class AgentSessionsToolWindowFactorySwingTest {
     val deleteTaskFolderIndex = entries.requiredIndex("AgentWorkbenchSessions.TreePopup.DeleteTaskFolder")
 
     assertThat(entries[newThreadIndex + 1]).isEqualTo("AgentWorkbenchSessions.TreePopup.CreateTaskFolder")
-    assertThat(entries[createTaskFolderIndex + 1]).isEqualTo(SEPARATOR_MARKER)
+    assertThat(entries[createTaskFolderIndex + 1]).isEqualTo("AgentWorkbenchSessions.TreePopup.TaskFolderAgent")
+    assertThat(entries[taskFolderAgentIndex + 1]).isEqualTo(SEPARATOR_MARKER)
     assertThat(moveToTaskFolderIndex).isLessThan(removeFromTaskFolderIndex)
     assertThat(removeFromTaskFolderIndex).isLessThan(markTaskFolderDoneIndex)
     assertThat(markTaskFolderDoneIndex).isLessThan(renameTaskFolderIndex)
diff --git a/plugins/agent-workbench/sessions/testSrc/AgentSessionArchiveServiceIntegrationTest.kt b/plugins/agent-workbench/sessions/testSrc/AgentSessionArchiveServiceIntegrationTest.kt
index e2c38a5f3271..e166be1ffe68 100644
--- a/plugins/agent-workbench/sessions/testSrc/AgentSessionArchiveServiceIntegrationTest.kt
+++ b/plugins/agent-workbench/sessions/testSrc/AgentSessionArchiveServiceIntegrationTest.kt
@@ -991,7 +991,13 @@ class AgentSessionArchiveServiceIntegrationTest {
 
           val threadToArchive = service.state.value.projects.first().threads.first { it.id == "codex-1" }
           val target = ArchiveThreadTarget.Thread(path = PROJECT_PATH, provider = threadToArchive.provider, threadId = threadToArchive.id)
-          archiveService.archiveThreadsForTest(listOf(target))
+          val archiveCompleted = CompletableDeferred<Unit>()
+          archiveService.archiveThreads(
+            targets = listOf(target),
+            entryPoint = AgentWorkbenchEntryPoint.TREE_POPUP,
+            onComplete = { archiveCompleted.complete(Unit) },
+          )
+          archiveCompleted.await()
           waitForCondition {
             service.state.value.projects.firstOrNull()?.threads?.none { it.id == "codex-1" } == true
           }