Optimistically render conversation items before subscription confirms

Author: charlesvien

apps/code/src/renderer/features/sessions/components/ConversationView.tsx

@@ -1,5 +1,6 @@
 import {
   sessionStoreSetters,
+  useOptimisticItemsForTask,
   usePendingPermissionsForTask,
   useQueuedMessagesForTask,
 } from "@features/sessions/stores/sessionStore";
@@ -68,6 +69,7 @@ export function ConversationView({
   const pendingPermissions = usePendingPermissionsForTask(taskId ?? "");
   const pendingPermissionsCount = pendingPermissions.size;
   const queuedMessages = useQueuedMessagesForTask(taskId);
+  const optimisticItems = useOptimisticItemsForTask(taskId);
 
   const queuedItems = useMemo<Extract<ConversationItem, { type: "queued" }>[]>(
     () =>
@@ -79,13 +81,10 @@ export function ConversationView({
     [queuedMessages],
   );
 
-  const items = useMemo<ConversationItem[]>(
-    () =>
-      queuedItems.length > 0
-        ? [...conversationItems, ...queuedItems]
-        : conversationItems,
-    [conversationItems, queuedItems],
-  );
+  const items = useMemo<ConversationItem[]>(() => {
+    const result: ConversationItem[] = [...conversationItems, ...optimisticItems];
+    return queuedItems.length > 0 ? [...result, ...queuedItems] : result;
+  }, [conversationItems, optimisticItems, queuedItems]);
 
   const handleScrollStateChange = useCallback((isAtBottom: boolean) => {
     setShowScrollButton(!isAtBottom);

apps/code/src/renderer/features/sessions/hooks/useSession.ts

@@ -10,6 +10,7 @@ import {
   type Adapter,
   type AgentSession,
   getConfigOptionByCategory,
+  type OptimisticItem,
   type QueuedMessage,
   useSessionStore,
 } from "../stores/sessionStore";
@@ -98,6 +99,17 @@ export const useQueuedMessagesForTask = (
   });
 };
 
+export const useOptimisticItemsForTask = (
+  taskId: string | undefined,
+): OptimisticItem[] => {
+  return useSessionStore((s) => {
+    if (!taskId) return [];
+    const taskRunId = s.taskIdIndex[taskId];
+    if (!taskRunId) return [];
+    return s.sessions[taskRunId]?.optimisticItems ?? [];
+  });
+};
+
 // --- Config Option Hooks ---
 
 /** Get a config option by category for a task */

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

@@ -819,6 +819,15 @@ export class SessionService {
     const session = sessionStoreSetters.getSessions()[taskRunId];
     if (!session) return;
 
+    if (
+      isJsonRpcRequest(acpMsg.message) &&
+      acpMsg.message.method === "session/prompt"
+    ) {
+      sessionStoreSetters.appendEventAndClearOptimisticItems(taskRunId, acpMsg);
+      this.updatePromptStateFromEvents(taskRunId, [acpMsg]);
+      return;
+    }
+
     sessionStoreSetters.appendEvents(taskRunId, [acpMsg]);
     this.updatePromptStateFromEvents(taskRunId, [acpMsg]);
 
@@ -1065,6 +1074,12 @@ export class SessionService {
       promptStartedAt: Date.now(),
     });
 
+    sessionStoreSetters.appendOptimisticItem(session.taskRunId, {
+      type: "user_message",
+      content: extractPromptText(blocks),
+      timestamp: Date.now(),
+    });
+
     try {
       const result = await trpcVanilla.agent.prompt.mutate({
         sessionId: session.taskRunId,
@@ -1082,6 +1097,8 @@ export class SessionService {
       const errorDetails = (error as { data?: { details?: string } }).data
         ?.details;
 
+      sessionStoreSetters.clearOptimisticItems(session.taskRunId);
+
       if (isFatalSessionError(errorMessage, errorDetails)) {
         log.error("Fatal prompt error, setting session to error state", {
           taskRunId: session.taskRunId,

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

@@ -28,6 +28,13 @@ export type TaskRunStatus =
   | "failed"
   | "cancelled";
 
+export type OptimisticItem = {
+  type: "user_message";
+  id: string;
+  content: string;
+  timestamp: number;
+};
+
 export interface AgentSession {
   taskRunId: string;
   taskId: string;
@@ -61,6 +68,7 @@ export interface AgentSession {
   cloudErrorMessage?: string | null;
   /** Cloud task branch */
   cloudBranch?: string | null;
+  optimisticItems?: OptimisticItem[];
 }
 
 // --- Config Option Helpers ---
@@ -189,6 +197,7 @@ export {
   useConfigOptionForTask,
   useModeConfigOptionForTask,
   useModelConfigOptionForTask,
+  useOptimisticItemsForTask,
   usePendingPermissionsForTask,
   useQueuedMessagesForTask,
   useSessionForTask,
@@ -332,6 +341,41 @@ export const sessionStoreSetters = {
     return result;
   },
 
+  appendOptimisticItem: (
+    taskRunId: string,
+    item: Omit<OptimisticItem, "id">,
+  ): void => {
+    const id = `optimistic-${Date.now()}-${Math.random().toString(36).slice(2, 7)}`;
+    useSessionStore.setState((state) => {
+      const session = state.sessions[taskRunId];
+      if (session) {
+        (session.optimisticItems ??= []).push({ ...item, id });
+      }
+    });
+  },
+
+  clearOptimisticItems: (taskRunId: string): void => {
+    useSessionStore.setState((state) => {
+      const session = state.sessions[taskRunId];
+      if (session) {
+        session.optimisticItems = [];
+      }
+    });
+  },
+
+  appendEventAndClearOptimisticItems: (
+    taskRunId: string,
+    event: AcpMessage,
+  ): void => {
+    useSessionStore.setState((state) => {
+      const session = state.sessions[taskRunId];
+      if (session) {
+        session.events.push(event);
+        session.optimisticItems = [];
+      }
+    });
+  },
+
   /** O(1) lookup using taskIdIndex */
   getSessionByTaskId: (taskId: string): AgentSession | undefined => {
     const state = useSessionStore.getState();