Handle session cleanup and error message status

packages/ui/src/stores/message-v2/instance-store.ts

@@ -182,7 +182,6 @@ function rebuildUsageStateFromInfos(infos: Iterable<MessageInfo>): SessionUsageS
 }
 
 export interface InstanceMessageStore {
-
   instanceId: string
   state: InstanceMessageState
   setState: SetStoreFunction<InstanceMessageState>
@@ -207,11 +206,14 @@ export interface InstanceMessageStore {
   getSessionRevision: (sessionId: string) => number
   getSessionMessageIds: (sessionId: string) => string[]
   getMessage: (messageId: string) => MessageRecord | undefined
+  clearSession: (sessionId: string) => void
   clearInstance: () => void
 }
 
 export function createInstanceMessageStore(instanceId: string): InstanceMessageStore {
   const [state, setState] = createStore<InstanceMessageState>(createInitialState(instanceId))
+
+
   const messageInfoCache = new Map<string, MessageInfo>()
 
   function bumpSessionRevision(sessionId: string) {
@@ -661,36 +663,109 @@ export function createInstanceMessageStore(instanceId: string): InstanceMessageS
     return state.scrollState[key]
   }
 
-  function clearInstance() {
-    messageInfoCache.clear()
-    setState(reconcile(createInitialState(instanceId)))
-  }
+   function clearSession(sessionId: string) {
+     if (!sessionId) return
+
+     const messageIds = Object.values(state.messages)
+       .filter((record) => record.sessionId === sessionId)
+       .map((record) => record.id)
+
+     // Remove message-level data
+     setState("messages", (prev) => {
+       const next = { ...prev }
+       messageIds.forEach((id) => delete next[id])
+       return next
+     })
+
+     setState("messageInfoVersion", (prev) => {
+       const next = { ...prev }
+       messageIds.forEach((id) => delete next[id])
+       return next
+     })
+
+     messageIds.forEach((id) => messageInfoCache.delete(id))
+
+     setState("pendingParts", (prev) => {
+       const next = { ...prev }
+       messageIds.forEach((id) => {
+         if (next[id]) delete next[id]
+       })
+       return next
+     })
+
+     setState("permissions", "byMessage", (prev) => {
+       const next = { ...prev }
+       messageIds.forEach((id) => {
+         if (next[id]) delete next[id]
+       })
+       return next
+     })
+
+     // Remove session-level data
+     setState("usage", (prev) => {
+       const next = { ...prev }
+       delete next[sessionId]
+       return next
+     })
+
+     setState("sessionRevisions", (prev) => {
+       const next = { ...prev }
+       delete next[sessionId]
+       return next
+     })
+
+     setState("scrollState", (prev) => {
+       const next = { ...prev }
+       const prefix = `${sessionId}:`
+       Object.keys(next).forEach((key) => {
+         if (key.startsWith(prefix)) {
+           delete next[key]
+         }
+       })
+       return next
+     })
+
+     setState("sessions", (prev) => {
+       const next = { ...prev }
+       delete next[sessionId]
+       return next
+     })
+
+     setState("sessionOrder", (ids) => ids.filter((id) => id !== sessionId))
+   }
+
+   function clearInstance() {
+     messageInfoCache.clear()
+     setState(reconcile(createInitialState(instanceId)))
+   }
+ 
+   return {
+     instanceId,
+     state,
+     setState,
+     addOrUpdateSession,
+     hydrateMessages,
+     upsertMessage,
+     applyPartUpdate,
+     bufferPendingPart,
+     flushPendingParts,
+     replaceMessageId,
+     setMessageInfo,
+     getMessageInfo,
+     upsertPermission,
+     removePermission,
+     getPermissionState,
+     setSessionRevert,
+     getSessionRevert,
+     rebuildUsage,
+     getSessionUsage,
+     setScrollSnapshot,
+     getScrollSnapshot,
+     getSessionRevision: getSessionRevisionValue,
+     getSessionMessageIds: (sessionId: string) => state.sessions[sessionId]?.messageIds ?? [],
+     getMessage: (messageId: string) => state.messages[messageId],
+     clearSession,
+     clearInstance,
+   }
+ }
 
-  return {
-    instanceId,
-    state,
-    setState,
-    addOrUpdateSession,
-    hydrateMessages,
-    upsertMessage,
-    applyPartUpdate,
-    bufferPendingPart,
-    flushPendingParts,
-    replaceMessageId,
-    setMessageInfo,
-    getMessageInfo,
-    upsertPermission,
-    removePermission,
-    getPermissionState,
-    setSessionRevert,
-    getSessionRevert,
-    rebuildUsage,
-    getSessionUsage,
-    setScrollSnapshot,
-    getScrollSnapshot,
-    getSessionRevision: getSessionRevisionValue,
-    getSessionMessageIds: (sessionId: string) => state.sessions[sessionId]?.messageIds ?? [],
-    getMessage: (messageId: string) => state.messages[messageId],
-    clearInstance,
-  }
-}

packages/ui/src/stores/session-api.ts

@@ -25,6 +25,8 @@ import { DEFAULT_MODEL_OUTPUT_LIMIT, getDefaultModel, isModelValid } from "./ses
 import { normalizeMessagePart } from "./message-v2/normalizers"
 import { updateSessionInfo } from "./message-v2/session-info"
 import { seedSessionMessagesV2 } from "./message-v2/bridge"
+import { messageStoreBus } from "./message-v2/bus"
+import { clearCacheForSession } from "../lib/global-cache"
 
 interface SessionForkResponse {
   id: string
@@ -358,6 +360,10 @@ async function deleteSession(instanceId: string, sessionId: string): Promise<voi
     setSessionCompactionState(instanceId, sessionId, false)
     clearSessionDraftPrompt(instanceId, sessionId)
 
+    // Drop normalized message state and caches for this session
+    messageStoreBus.getOrCreate(instanceId).clearSession(sessionId)
+    clearCacheForSession(instanceId, sessionId)
+
     setSessionInfoByInstance((prev) => {
       const next = new Map(prev)
       const instanceInfo = next.get(instanceId)

packages/ui/src/stores/session-events.ts

@@ -13,6 +13,7 @@ import type {
   EventSessionIdle,
   EventSessionUpdated,
 } from "@opencode-ai/sdk"
+import type { MessageStatus } from "./message-v2/types"
 
 import { showToastNotification, ToastVariant } from "../lib/notifications"
 import { instances, addPermissionToQueue, removePermissionFromQueue, refreshPermissionsForSession } from "./instances"
@@ -136,6 +137,8 @@ function handleMessageUpdate(instanceId: string, event: MessageUpdateEvent | Mes
 
     const store = messageStoreBus.getOrCreate(instanceId)
     const role: MessageRole = info.role === "user" ? "user" : "assistant"
+    const hasError = Boolean((info as any).error)
+    const status: MessageStatus = hasError ? "error" : "complete"
 
     let record = store.getMessage(messageId)
     if (!record) {
@@ -153,18 +156,19 @@ function handleMessageUpdate(instanceId: string, event: MessageUpdateEvent | Mes
         id: messageId,
         sessionId,
         role,
-        status: "complete",
+        status,
         createdAt,
         updatedAt: completedAt ?? createdAt,
       })
     }
 
-    upsertMessageInfoV2(instanceId, info, { status: "complete", bumpRevision: true })
+    upsertMessageInfoV2(instanceId, info, { status, bumpRevision: true })
 
     updateSessionInfo(instanceId, sessionId)
     refreshPermissionsForSession(instanceId, sessionId)
   }
-}
+ }
+
 
 function handleSessionUpdate(instanceId: string, event: EventSessionUpdated): void {
   const info = event.properties?.info