From 6a9a442948c9122478fac92f3d86344a5c436c32 Mon Sep 17 00:00:00 2001 From: Shantur Rathore Date: Wed, 26 Nov 2025 16:20:02 +0000 Subject: [PATCH] Handle session cleanup and error message status --- .../src/stores/message-v2/instance-store.ts | 141 ++++++++++++++---- packages/ui/src/stores/session-api.ts | 6 + packages/ui/src/stores/session-events.ts | 10 +- 3 files changed, 121 insertions(+), 36 deletions(-) diff --git a/packages/ui/src/stores/message-v2/instance-store.ts b/packages/ui/src/stores/message-v2/instance-store.ts index f4386b11..33cac7d6 100644 --- a/packages/ui/src/stores/message-v2/instance-store.ts +++ b/packages/ui/src/stores/message-v2/instance-store.ts @@ -182,7 +182,6 @@ function rebuildUsageStateFromInfos(infos: Iterable): SessionUsageS } export interface InstanceMessageStore { - instanceId: string state: InstanceMessageState setState: SetStoreFunction @@ -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(createInitialState(instanceId)) + + const messageInfoCache = new Map() 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, - } -} diff --git a/packages/ui/src/stores/session-api.ts b/packages/ui/src/stores/session-api.ts index 5337c7ce..4d00cee8 100644 --- a/packages/ui/src/stores/session-api.ts +++ b/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 { const next = new Map(prev) const instanceInfo = next.get(instanceId) diff --git a/packages/ui/src/stores/session-events.ts b/packages/ui/src/stores/session-events.ts index 0eb42233..5c6962ff 100644 --- a/packages/ui/src/stores/session-events.ts +++ b/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