From 2603b1d260c7fc4b02d0199e98921b4c18e10383 Mon Sep 17 00:00:00 2001 From: Shantur Rathore Date: Thu, 25 Dec 2025 15:12:44 +0000 Subject: [PATCH] Handle revert removals locally and retarget prompt input --- .gitignore | 3 +- .../src/components/session/session-view.tsx | 7 +- .../ui/src/lib/hooks/use-app-lifecycle.ts | 8 +- packages/ui/src/lib/hooks/use-commands.ts | 20 +++- packages/ui/src/stores/message-v2/bridge.ts | 12 ++ .../src/stores/message-v2/instance-store.ts | 110 +++++++++++++++++- packages/ui/src/stores/session-events.ts | 20 ++-- 7 files changed, 158 insertions(+), 22 deletions(-) diff --git a/.gitignore b/.gitignore index 39636667..da104bdd 100644 --- a/.gitignore +++ b/.gitignore @@ -6,4 +6,5 @@ release/ .vite/ .electron-vite/ out/ -.dir-locals.el \ No newline at end of file +.dir-locals.el +.opencode/bashOutputs/ \ No newline at end of file diff --git a/packages/ui/src/components/session/session-view.tsx b/packages/ui/src/components/session/session-view.tsx index 9e701e55..19d21a00 100644 --- a/packages/ui/src/components/session/session-view.tsx +++ b/packages/ui/src/components/session/session-view.tsx @@ -39,6 +39,7 @@ export const SessionView: Component = (props) => { return getSessionBusyStatus(props.instanceId, currentSession.id) }) let scrollToBottomHandle: (() => void) | undefined + let rootRef: HTMLDivElement | undefined function scheduleScrollToBottom() { if (!scrollToBottomHandle) return requestAnimationFrame(() => { @@ -128,7 +129,7 @@ export const SessionView: Component = (props) => { const restoredText = getUserMessageText(messageId) if (restoredText) { - const textarea = document.querySelector(".prompt-input") as HTMLTextAreaElement + const textarea = rootRef?.querySelector(".prompt-input") as HTMLTextAreaElement | undefined if (textarea) { textarea.value = restoredText textarea.dispatchEvent(new Event("input", { bubbles: true })) @@ -164,7 +165,7 @@ export const SessionView: Component = (props) => { await loadMessages(props.instanceId, forkedSession.id).catch((error) => log.error("Failed to load forked session messages", error)) if (restoredText) { - const textarea = document.querySelector(".prompt-input") as HTMLTextAreaElement + const textarea = rootRef?.querySelector(".prompt-input") as HTMLTextAreaElement | undefined if (textarea) { textarea.value = restoredText textarea.dispatchEvent(new Event("input", { bubbles: true })) @@ -194,7 +195,7 @@ export const SessionView: Component = (props) => { const activeSession = sessionAccessor() if (!activeSession) return null return ( -
+
{ - const textarea = document.querySelector(".prompt-input") as HTMLTextAreaElement + const textarea = document.querySelector( + ".session-cache-pane[aria-hidden=\"false\"] .prompt-input", + ) as HTMLTextAreaElement if (textarea) textarea.value = "" }, () => { - const textarea = document.querySelector(".prompt-input") as HTMLTextAreaElement + const textarea = document.querySelector( + ".session-cache-pane[aria-hidden=\"false\"] .prompt-input", + ) as HTMLTextAreaElement textarea?.focus() }, ) diff --git a/packages/ui/src/lib/hooks/use-commands.ts b/packages/ui/src/lib/hooks/use-commands.ts index 7c2eedcb..d6155f55 100644 --- a/packages/ui/src/lib/hooks/use-commands.ts +++ b/packages/ui/src/lib/hooks/use-commands.ts @@ -261,6 +261,22 @@ export function useCommands(options: UseCommandsOptions) { }, }) + function escapeCss(value: string) { + if (typeof CSS !== "undefined" && typeof (CSS as any).escape === "function") { + return (CSS as any).escape(value) + } + return value.replace(/\\/g, "\\\\").replace(/"/g, "\\\"") + } + + function findVisiblePromptTextarea(sessionId?: string): HTMLTextAreaElement | null { + if (typeof document === "undefined") return null + const base = ".session-cache-pane[aria-hidden=\"false\"]" + const selector = sessionId + ? `${base}[data-session-id=\"${escapeCss(sessionId)}\"] .prompt-input` + : `${base} .prompt-input` + return document.querySelector(selector) as HTMLTextAreaElement | null + } + commandRegistry.register({ id: "undo", label: "Undo Last Message", @@ -327,7 +343,7 @@ export function useCommands(options: UseCommandsOptions) { } if (restoredText) { - const textarea = document.querySelector(".prompt-input") as HTMLTextAreaElement + const textarea = findVisiblePromptTextarea(sessionId) if (textarea) { textarea.value = restoredText textarea.dispatchEvent(new Event("input", { bubbles: true })) @@ -381,7 +397,7 @@ export function useCommands(options: UseCommandsOptions) { keywords: ["clear", "reset"], shortcut: { key: "K", meta: true }, action: () => { - const textarea = document.querySelector(".prompt-input") as HTMLTextAreaElement + const textarea = findVisiblePromptTextarea() if (textarea) textarea.value = "" }, }) diff --git a/packages/ui/src/stores/message-v2/bridge.ts b/packages/ui/src/stores/message-v2/bridge.ts index 2226eeba..5252e0f6 100644 --- a/packages/ui/src/stores/message-v2/bridge.ts +++ b/packages/ui/src/stores/message-v2/bridge.ts @@ -143,6 +143,18 @@ export function removePermissionV2(instanceId: string, permissionId: string): vo store.removePermission(permissionId) } +export function removeMessageV2(instanceId: string, messageId: string): void { + if (!messageId) return + const store = messageStoreBus.getOrCreate(instanceId) + store.removeMessage(messageId) +} + +export function removeMessagePartV2(instanceId: string, messageId: string, partId: string): void { + if (!messageId || !partId) return + const store = messageStoreBus.getOrCreate(instanceId) + store.removeMessagePart(messageId, partId) +} + export function ensureSessionMetadataV2(instanceId: string, session: Session | null | undefined): void { if (!session) return const store = messageStoreBus.getOrCreate(instanceId) diff --git a/packages/ui/src/stores/message-v2/instance-store.ts b/packages/ui/src/stores/message-v2/instance-store.ts index cf6ad570..c51eb1f6 100644 --- a/packages/ui/src/stores/message-v2/instance-store.ts +++ b/packages/ui/src/stores/message-v2/instance-store.ts @@ -191,6 +191,8 @@ export interface InstanceMessageStore { hydrateMessages: (sessionId: string, inputs: MessageUpsertInput[], infos?: Iterable) => void upsertMessage: (input: MessageUpsertInput) => void applyPartUpdate: (input: PartUpdateInput) => void + removeMessage: (messageId: string) => void + removeMessagePart: (messageId: string, partId: string) => void bufferPendingPart: (entry: PendingPartEntry) => void flushPendingParts: (messageId: string) => void replaceMessageId: (options: ReplaceMessageIdOptions) => void @@ -508,10 +510,10 @@ export function createInstanceMessageStore(instanceId: string, hooks?: MessageSt bufferPendingPart({ messageId: input.messageId, part: input.part, receivedAt: Date.now() }) return } - + const partId = ensurePartId(input.messageId, input.part, message.partIds.length) const cloned = clonePart(input.part) - + setState( "messages", input.messageId, @@ -520,7 +522,7 @@ export function createInstanceMessageStore(instanceId: string, hooks?: MessageSt draft.partIds = [...draft.partIds, partId] } const existing = draft.parts[partId] - const nextRevision = existing ? existing.revision + 1 : cloned.version ?? 0 + const nextRevision = existing ? existing.revision + 1 : (cloned as any).version ?? 0 draft.parts[partId] = { id: partId, data: cloned, @@ -540,12 +542,106 @@ export function createInstanceMessageStore(instanceId: string, hooks?: MessageSt timestamp: Date.now(), }) } - + // Any part update can change the rendered height of the message // list, so we treat it as a session revision for scroll purposes. bumpSessionRevision(message.sessionId) } + function removeMessage(messageId: string) { + if (!messageId) return + + const record = state.messages[messageId] + const sessionIds = new Set() + + if (record?.sessionId) { + sessionIds.add(record.sessionId) + } else { + Object.values(state.sessions).forEach((session) => { + if (session.messageIds.includes(messageId)) { + sessionIds.add(session.id) + } + }) + } + + clearRecordDisplayCacheForMessages(instanceId, [messageId]) + + batch(() => { + sessionIds.forEach((sessionId) => { + setState("sessions", sessionId, "messageIds", (ids = []) => ids.filter((id) => id !== messageId)) + }) + + setState("messages", (prev) => { + if (!prev[messageId]) return prev + const next = { ...prev } + delete next[messageId] + return next + }) + + setState("messageInfoVersion", (prev) => { + if (!(messageId in prev)) return prev + const next = { ...prev } + delete next[messageId] + return next + }) + + messageInfoCache.delete(messageId) + + setState("pendingParts", (prev) => { + if (!prev[messageId]) return prev + const next = { ...prev } + delete next[messageId] + return next + }) + + setState("permissions", "byMessage", (prev) => { + if (!prev[messageId]) return prev + const next = { ...prev } + delete next[messageId] + return next + }) + + sessionIds.forEach((sessionId) => { + withUsageState(sessionId, (draft) => removeUsageEntry(draft, messageId)) + if (state.latestTodos[sessionId]?.messageId === messageId) { + clearLatestTodoSnapshot(sessionId) + } + bumpSessionRevision(sessionId) + }) + }) + } + + function removeMessagePart(messageId: string, partId: string) { + if (!messageId || !partId) return + const message = state.messages[messageId] + if (!message) return + + clearRecordDisplayCacheForMessages(instanceId, [messageId]) + + batch(() => { + setState( + "messages", + messageId, + produce((draft: MessageRecord) => { + if (!draft.parts[partId] && !draft.partIds.includes(partId)) return + draft.partIds = draft.partIds.filter((id) => id !== partId) + delete draft.parts[partId] + draft.updatedAt = Date.now() + draft.revision += 1 + }), + ) + + setState("permissions", "byMessage", messageId, (prev) => { + if (!prev || !prev[partId]) return prev + const next = { ...prev } + delete next[partId] + return next + }) + + bumpSessionRevision(message.sessionId) + }) + } + function flushPendingParts(messageId: string) { const pending = state.pendingParts[messageId] @@ -868,8 +964,10 @@ export function createInstanceMessageStore(instanceId: string, hooks?: MessageSt addOrUpdateSession, hydrateMessages, upsertMessage, - applyPartUpdate, - bufferPendingPart, + applyPartUpdate, + removeMessage, + removeMessagePart, + bufferPendingPart, flushPendingParts, replaceMessageId, setMessageInfo, diff --git a/packages/ui/src/stores/session-events.ts b/packages/ui/src/stores/session-events.ts index edbbe5db..9d13e055 100644 --- a/packages/ui/src/stores/session-events.ts +++ b/packages/ui/src/stores/session-events.ts @@ -31,6 +31,8 @@ import { replaceMessageIdV2, upsertMessageInfoV2, upsertPermissionV2, + removeMessagePartV2, + removeMessageV2, removePermissionV2, setSessionRevertV2, } from "./message-v2/bridge" @@ -305,19 +307,21 @@ function handleSessionError(_instanceId: string, event: EventSessionError): void } function handleMessageRemoved(instanceId: string, event: MessageRemovedEvent): void { - const sessionID = event.properties?.sessionID - if (!sessionID) return + const { sessionID, messageID } = event.properties + if (!sessionID || !messageID) return - log.info(`[SSE] Message removed from session ${sessionID}, reloading messages`) - loadMessages(instanceId, sessionID, true).catch((error) => log.error("Failed to reload messages after removal", error)) + log.info(`[SSE] Message removed from session ${sessionID}`, { messageID }) + removeMessageV2(instanceId, messageID) + updateSessionInfo(instanceId, sessionID) } function handleMessagePartRemoved(instanceId: string, event: MessagePartRemovedEvent): void { - const sessionID = event.properties?.sessionID - if (!sessionID) return + const { sessionID, messageID, partID } = event.properties + if (!sessionID || !messageID || !partID) return - log.info(`[SSE] Message part removed from session ${sessionID}, reloading messages`) - loadMessages(instanceId, sessionID, true).catch((error) => log.error("Failed to reload messages after part removal", error)) + log.info(`[SSE] Message part removed from session ${sessionID}`, { messageID, partID }) + removeMessagePartV2(instanceId, messageID, partID) + updateSessionInfo(instanceId, sessionID) } function handleTuiToast(_instanceId: string, event: TuiToastEvent): void {