Fix optimistic items review issues

Author: charlesvien

apps/code/src/renderer/features/sessions/service/service.test.ts

@@ -52,6 +52,7 @@ const mockSessionStoreSetters = vi.hoisted(() => ({
   clearAll: vi.fn(),
   appendOptimisticItem: vi.fn(),
   clearOptimisticItems: vi.fn(),
+  replaceOptimisticWithEvent: vi.fn(),
 }));
 
 vi.mock("@features/sessions/stores/sessionStore", () => ({
@@ -207,6 +208,7 @@ const createMockSession = (
   promptStartedAt: null,
   pendingPermissions: new Map(),
   messageQueue: [],
+  optimisticItems: [],
   ...overrides,
 });
 

apps/code/src/renderer/features/sessions/service/service.ts

@@ -817,16 +817,15 @@ export class SessionService {
     const session = sessionStoreSetters.getSessions()[taskRunId];
     if (!session) return;
 
-    if (
+    const isUserPromptEcho =
       isJsonRpcRequest(acpMsg.message) &&
-      acpMsg.message.method === "session/prompt"
-    ) {
-      sessionStoreSetters.appendEventAndClearOptimisticItems(taskRunId, acpMsg);
-      this.updatePromptStateFromEvents(taskRunId, [acpMsg]);
-      return;
-    }
+      acpMsg.message.method === "session/prompt";
 
-    sessionStoreSetters.appendEvents(taskRunId, [acpMsg]);
+    if (isUserPromptEcho) {
+      sessionStoreSetters.replaceOptimisticWithEvent(taskRunId, acpMsg);
+    } else {
+      sessionStoreSetters.appendEvents(taskRunId, [acpMsg]);
+    }
     this.updatePromptStateFromEvents(taskRunId, [acpMsg]);
 
     const msg = acpMsg.message;
@@ -1004,7 +1003,7 @@ export class SessionService {
       prompt_length_chars: promptText.length,
     });
 
-    return this.sendLocalPrompt(session, blocks);
+    return this.sendLocalPrompt(session, blocks, promptText);
   }
 
   /**
@@ -1051,7 +1050,7 @@ export class SessionService {
     });
 
     try {
-      return await this.sendLocalPrompt(session, blocks);
+      return await this.sendLocalPrompt(session, blocks, combinedText);
     } catch (error) {
       // Log that queued messages were lost due to send failure
       log.error("Failed to send queued messages, messages lost", {
@@ -1066,6 +1065,7 @@ export class SessionService {
   private async sendLocalPrompt(
     session: AgentSession,
     blocks: ContentBlock[],
+    promptText: string,
   ): Promise<{ stopReason: string }> {
     sessionStoreSetters.updateSession(session.taskRunId, {
       isPromptPending: true,
@@ -1074,7 +1074,7 @@ export class SessionService {
 
     sessionStoreSetters.appendOptimisticItem(session.taskRunId, {
       type: "user_message",
-      content: extractPromptText(blocks),
+      content: promptText,
       timestamp: Date.now(),
     });
 
@@ -1083,7 +1083,6 @@ export class SessionService {
         sessionId: session.taskRunId,
         prompt: blocks,
       });
-      // Clear pending state on success
       sessionStoreSetters.updateSession(session.taskRunId, {
         isPromptPending: false,
         promptStartedAt: null,
@@ -1955,6 +1954,7 @@ export class SessionService {
       promptStartedAt: null,
       pendingPermissions: new Map(),
       messageQueue: [],
+      optimisticItems: [],
     };
   }
 

apps/code/src/renderer/features/sessions/stores/sessionStore.ts

@@ -68,7 +68,7 @@ export interface AgentSession {
   cloudErrorMessage?: string | null;
   /** Cloud task branch */
   cloudBranch?: string | null;
-  optimisticItems?: OptimisticItem[];
+  optimisticItems: OptimisticItem[];
 }
 
 // --- Config Option Helpers ---
@@ -345,13 +345,10 @@ export const sessionStoreSetters = {
     taskRunId: string,
     item: Omit<OptimisticItem, "id">,
   ): void => {
-    const id = `optimistic-${Date.now()}-${Math.random().toString(36).slice(2, 7)}`;
+    const id = `optimistic-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`;
     useSessionStore.setState((state) => {
       const session = state.sessions[taskRunId];
       if (session) {
-        if (!session.optimisticItems) {
-          session.optimisticItems = [];
-        }
         session.optimisticItems.push({ ...item, id });
       }
     });
@@ -366,10 +363,7 @@ export const sessionStoreSetters = {
     });
   },
 
-  appendEventAndClearOptimisticItems: (
-    taskRunId: string,
-    event: AcpMessage,
-  ): void => {
+  replaceOptimisticWithEvent: (taskRunId: string, event: AcpMessage): void => {
     useSessionStore.setState((state) => {
       const session = state.sessions[taskRunId];
       if (session) {