fix(ui): keep resumed prompt before restore progress

Author: tatoalo

packages/ui/src/features/sessions/components/mergeConversationItems.test.ts

@@ -2,6 +2,24 @@ import { describe, expect, it } from "vitest";
 import type { ConversationItem } from "./buildConversationItems";
 import { mergeConversationItems } from "./mergeConversationItems";
 
+function progressGroup(id: string): ConversationItem {
+  return {
+    type: "session_update",
+    id,
+    update: {
+      sessionUpdate: "progress_group",
+      steps: [],
+      isActive: true,
+    },
+    turnContext: {
+      toolCalls: new Map(),
+      childItems: new Map(),
+      turnCancelled: false,
+      turnComplete: false,
+    },
+  };
+}
+
 function userMessage(
   id: string,
   content: string,
@@ -160,6 +178,19 @@ describe("mergeConversationItems", () => {
     expect(result.map((i) => i.id)).toEqual(["setup", "opt"]);
   });
 
+  it("cloud: keeps a resumed prompt before trailing setup progress", () => {
+    const result = mergeConversationItems({
+      conversationItems: [
+        userMessage("old", "previous prompt"),
+        progressGroup("restore"),
+      ],
+      optimisticItems: [userMessage("opt", "follow up", false)],
+      isCloud: true,
+    });
+
+    expect(result.map((item) => item.id)).toEqual(["old", "opt", "restore"]);
+  });
+
   it("cloud: does not dedupe historical messages against tail follow-up optimistics", () => {
     const result = mergeConversationItems({
       conversationItems: [

packages/ui/src/features/sessions/components/mergeConversationItems.ts

@@ -28,8 +28,9 @@ function strippedUserContent(content: string): string {
 }
 
 // Cloud's initial optimistic is pinned to the top so the user's prompt stays
-// visible above setup progress. Follow-up optimistics render at the tail until
-// the streamed `session/prompt` arrives and replaces them.
+// visible above setup progress. Follow-up optimistics render at the tail, but
+// before trailing progress cards, to match where the streamed `session/prompt`
+// will appear.
 //
 // Local sessions keep optimistic at the chronological end — they rely on
 // `replaceOptimisticWithEvent` to swap optimistic↔real in place.
@@ -103,9 +104,23 @@ export function mergeConversationItems({
           };
         });
 
+  let tailInsertionIndex = dedupedConversation.length;
+  while (tailInsertionIndex > 0) {
+    const item = dedupedConversation[tailInsertionIndex - 1];
+    if (
+      item.type === "session_update" &&
+      item.update.sessionUpdate === "progress_group"
+    ) {
+      tailInsertionIndex--;
+    } else {
+      break;
+    }
+  }
+
   return [
     ...resolvedPinnedItems,
-    ...dedupedConversation,
+    ...dedupedConversation.slice(0, tailInsertionIndex),
     ...tailOptimisticItems,
+    ...dedupedConversation.slice(tailInsertionIndex),
   ];
 }

packages/ui/src/features/sessions/sessionServiceHost.test.ts

@@ -356,6 +356,9 @@ const mockConvertStoredEntriesToEvents = vi.hoisted(() =>
     ) => unknown[]
   >(() => []),
 );
+const mockHasSessionPromptEventForTaskRun = vi.hoisted(() =>
+  vi.fn(() => false),
+);
 
 vi.mock("@posthog/core/sessions/sessionEvents", async () => {
   const actual = await vi.importActual<
@@ -387,6 +390,7 @@ vi.mock("@posthog/core/sessions/sessionEvents", async () => {
     getStoredLogEventPosition: actual.getStoredLogEventPosition,
     getUserShellExecutesSinceLastPrompt: vi.fn(() => []),
     hasSessionPromptEvent: actual.hasSessionPromptEvent,
+    hasSessionPromptEventForTaskRun: mockHasSessionPromptEventForTaskRun,
     isAbsoluteFolderPath: actual.isAbsoluteFolderPath,
     isFatalSessionError: actual.isFatalSessionError,
     isRateLimitError: actual.isRateLimitError,
@@ -446,6 +450,7 @@ describe("SessionService", () => {
   beforeEach(() => {
     vi.clearAllMocks();
     mockConvertStoredEntriesToEvents.mockImplementation(() => []);
+    mockHasSessionPromptEventForTaskRun.mockReturnValue(false);
     resetSessionService();
     mockSettingsState.customInstructions = "";
     mockSettingsState.spokenNotifications = false;
@@ -4374,6 +4379,7 @@ describe("SessionService", () => {
           resumePrompt,
           resumeCompletion,
         ]);
+        mockHasSessionPromptEventForTaskRun.mockReturnValueOnce(true);
 
         service.watchCloudTask(
           "task-123",