diff --git a/packages/ui/src/components/message-stream-v2.tsx b/packages/ui/src/components/message-stream-v2.tsx new file mode 100644 index 00000000..44e8fd26 --- /dev/null +++ b/packages/ui/src/components/message-stream-v2.tsx @@ -0,0 +1,422 @@ +import { For, Show, createMemo, createSignal, createEffect, onCleanup } from "solid-js" +import MessageItem from "./message-item" +import ToolCall from "./tool-call" +import Kbd from "./kbd" +import type { Message, MessageInfo, ClientPart } from "../types/message" +import { computeDisplayParts } from "../stores/session-messages" +import { getSessionInfo } from "../stores/sessions" +import { showCommandPalette } from "../stores/command-palette" +import { messageStoreBus } from "../stores/message-v2/bus" +import type { MessageRecord } from "../stores/message-v2/types" +import { useConfig } from "../stores/preferences" +import { getScrollCache, setScrollCache } from "../lib/scroll-cache" +import { sseManager } from "../lib/sse-manager" +import { formatTokenTotal } from "../lib/formatters" + +const SCROLL_SCOPE = "session" +const TOOL_ICON = "🔧" +const codeNomadLogo = new URL("../images/CodeNomad-Icon.png", import.meta.url).href + +function formatTokens(tokens: number): string { + return formatTokenTotal(tokens) +} + + +interface MessageStreamV2Props { + instanceId: string + sessionId: string + loading?: boolean + onRevert?: (messageId: string) => void + onFork?: (messageId?: string) => void +} + +interface MessageDisplayItem { + type: "message" + message: Message + combinedParts: ClientPart[] + messageInfo?: MessageInfo + isQueued: boolean +} + +interface ToolDisplayItem { + type: "tool" + key: string + toolPart: ToolCallPart + messageInfo?: MessageInfo + messageId: string + messageVersion: number + partVersion: number +} + +type DisplayItem = MessageDisplayItem | ToolDisplayItem + +type ToolCallPart = Extract + +function isToolPart(part: ClientPart): part is ToolCallPart { + return part.type === "tool" +} + +function recordToMessage(record: MessageRecord): Message { + const parts = record.partIds + .map((partId) => record.parts[partId]?.data) + .filter((part): part is ClientPart => Boolean(part)) + return { + id: record.id, + sessionId: record.sessionId, + type: record.role, + parts, + timestamp: record.createdAt, + status: record.status, + version: record.revision, + } +} + +function hasRenderableContent(message: Message, combinedParts: ClientPart[], info?: MessageInfo): boolean { + if (message.type !== "assistant" && message.type !== "user") { + return false + } + if (message.type !== "assistant" || combinedParts.length > 0) { + return true + } + if (info && info.role === "assistant" && info.error) { + return true + } + return message.status === "error" +} + +export default function MessageStreamV2(props: MessageStreamV2Props) { + const { preferences } = useConfig() + const store = createMemo(() => messageStoreBus.getOrCreate(props.instanceId)) + const messageIds = createMemo(() => store().getSessionMessageIds(props.sessionId)) + const messageRecords = createMemo(() => + messageIds() + .map((id) => store().getMessage(id)) + .filter((record): record is MessageRecord => Boolean(record)), + ) + + const usageSnapshot = createMemo(() => store().getSessionUsage(props.sessionId)) + const sessionInfo = createMemo(() => + getSessionInfo(props.instanceId, props.sessionId) ?? { + cost: 0, + contextWindow: 0, + isSubscriptionModel: false, + inputTokens: 0, + outputTokens: 0, + reasoningTokens: 0, + actualUsageTokens: 0, + modelOutputLimit: 0, + contextAvailableTokens: null, + }, + ) + const tokenStats = createMemo(() => { + const usage = usageSnapshot() + const info = sessionInfo() + return { + used: usage?.actualUsageTokens ?? info.actualUsageTokens ?? 0, + avail: info.contextAvailableTokens, + } + }) + + const connectionStatus = () => sseManager.getStatus(props.instanceId) + const handleCommandPaletteClick = () => { + showCommandPalette(props.instanceId) + } + + const messageInfoMap = createMemo(() => { + const map = new Map() + messageIds().forEach((id) => { + const info = store().getMessageInfo(id) + if (info) { + map.set(id, info) + } + }) + return map + }) + const revertTarget = createMemo(() => store().getSessionRevert(props.sessionId)) + + const displayItems = createMemo(() => { + const infoMap = messageInfoMap() + const showThinking = preferences().showThinkingBlocks + const revert = revertTarget() + const items: DisplayItem[] = [] + + const records = messageRecords() + let lastAssistantIndex = -1 + for (let i = records.length - 1; i >= 0; i--) { + if (records[i].role === "assistant") { + lastAssistantIndex = i + break + } + } + + for (let index = 0; index < records.length; index++) { + const record = records[index] + if (revert?.messageID && record.id === revert.messageID) { + break + } + + const baseMessage = recordToMessage(record) + const displayParts = computeDisplayParts(baseMessage, showThinking) + baseMessage.displayParts = displayParts + const combinedParts = displayParts.combined + const messageInfo = infoMap.get(record.id) + const isQueued = + baseMessage.type === "user" && (lastAssistantIndex === -1 || index > lastAssistantIndex) + + if (hasRenderableContent(baseMessage, combinedParts, messageInfo)) { + items.push({ + type: "message", + message: baseMessage, + combinedParts, + messageInfo, + isQueued, + }) + } + + const toolParts: ToolCallPart[] = displayParts.tool.filter(isToolPart) + toolParts.forEach((toolPart, toolIndex) => { + const partVersion = typeof toolPart.version === "number" ? toolPart.version : 0 + const messageVersion = typeof baseMessage.version === "number" ? baseMessage.version : 0 + const key = toolPart.id || `${record.id}-tool-${toolIndex}` + items.push({ + type: "tool", + key, + toolPart, + messageInfo, + messageId: record.id, + messageVersion, + partVersion, + }) + }) + } + + return items + }) + + const changeToken = createMemo(() => { + const entries = displayItems() + return entries + .map((item) => { + if (item.type === "message") { + return `${item.message.id}:${item.message.version}:${item.combinedParts.length}` + } + const status = item.toolPart.state?.status || "unknown" + return `tool:${item.key}:${item.partVersion}:${status}` + }) + .join("|") + }) + + const [autoScroll, setAutoScroll] = createSignal(true) + const [showScrollButton, setShowScrollButton] = createSignal(false) + let containerRef: HTMLDivElement | undefined + + function isNearBottom(element: HTMLDivElement, offset = 48) { + const { scrollTop, scrollHeight, clientHeight } = element + return scrollHeight - (scrollTop + clientHeight) <= offset + } + + function scrollToBottom(immediate = false) { + if (!containerRef) return + const behavior = immediate ? "auto" : "smooth" + containerRef.scrollTo({ top: containerRef.scrollHeight, behavior }) + setAutoScroll(true) + persistScrollState() + } + + function persistScrollState() { + if (!containerRef) return + setScrollCache( + { instanceId: props.instanceId, sessionId: props.sessionId, scope: SCROLL_SCOPE }, + { + scrollTop: containerRef.scrollTop, + atBottom: isNearBottom(containerRef), + }, + ) + } + + function handleScroll(event: Event) { + if (!containerRef) return + const atBottom = isNearBottom(containerRef) + setShowScrollButton(!atBottom) + if (event.isTrusted) { + setAutoScroll(atBottom) + } + persistScrollState() + } + + createEffect(() => { + const scrollSnapshot = getScrollCache({ instanceId: props.instanceId, sessionId: props.sessionId, scope: SCROLL_SCOPE }) + requestAnimationFrame(() => { + if (!containerRef) return + if (scrollSnapshot) { + const maxScrollTop = Math.max(containerRef.scrollHeight - containerRef.clientHeight, 0) + containerRef.scrollTop = Math.min(scrollSnapshot.scrollTop, maxScrollTop) + setAutoScroll(scrollSnapshot.atBottom) + setShowScrollButton(!scrollSnapshot.atBottom) + } else { + scrollToBottom(true) + } + }) + }) + + let previousToken: string | undefined + createEffect(() => { + const token = changeToken() + if (!token || token === previousToken) { + return + } + previousToken = token + if (autoScroll()) { + requestAnimationFrame(() => scrollToBottom(true)) + } + }) + + createEffect(() => { + if (displayItems().length === 0) { + setShowScrollButton(false) + setAutoScroll(true) + } + }) + + onCleanup(() => { + persistScrollState() + }) + + return ( +
+
+
+
+ Used + {formatTokens(tokenStats().used)} +
+
+ Avail + + {sessionInfo().contextAvailableTokens !== null ? formatTokens(sessionInfo().contextAvailableTokens ?? 0) : "--"} + +
+
+ +
+
+ + + + +
+
+
+ + + + Connected + + + + + + Connecting... + + + + + + Disconnected + + +
+
+ +
{ + containerRef = element || undefined + }} + onScroll={handleScroll} + > + +
+
+
+ CodeNomad logo +

CodeNomad

+
+

Start a conversation

+

Type a message below or open the Command Palette:

+
    +
  • + Command Palette + +
  • +
  • Ask about your codebase
  • +
  • + Attach files with @ +
  • +
+
+
+
+ + +
+
+

Loading messages...

+
+ + + + {(item) => { + if (item.type === "message") { + return ( + + ) + } + + return ( +
+
+
+ {TOOL_ICON} + Tool Call + {item.toolPart.tool || "unknown"} +
+
+ +
+ ) + }} +
+
+ + +
+ +
+
+
+ ) +} diff --git a/packages/ui/src/components/session/session-view.tsx b/packages/ui/src/components/session/session-view.tsx index 04d35b54..dd269be0 100644 --- a/packages/ui/src/components/session/session-view.tsx +++ b/packages/ui/src/components/session/session-view.tsx @@ -2,12 +2,17 @@ import { Show, createMemo, createEffect, onCleanup, type Component } from "solid import type { Session } from "../../types/session" import type { Attachment } from "../../types/attachment" import type { ClientPart } from "../../types/message" -import MessageStream from "../message-stream" +import MessageStreamV2 from "../message-stream-v2" +import { messageStoreBus } from "../../stores/message-v2/bus" import PromptInput from "../prompt-input" import { instances } from "../../stores/instances" import { loadMessages, sendMessage, forkSession, isSessionMessagesLoading, setActiveParentSession, setActiveSession, runShellCommand } from "../../stores/sessions" import { showAlertDialog } from "../../stores/alerts" +function isTextPart(part: ClientPart): part is ClientPart & { type: "text"; text: string } { + return part?.type === "text" && typeof (part as any).text === "string" +} + interface SessionViewProps { sessionId: string activeSessions: Map @@ -19,6 +24,7 @@ interface SessionViewProps { export const SessionView: Component = (props) => { const session = () => props.activeSessions.get(props.sessionId) const messagesLoading = createMemo(() => isSessionMessagesLoading(props.instanceId, props.sessionId)) + const messageStore = createMemo(() => messageStoreBus.getOrCreate(props.instanceId)) createEffect(() => { const currentSession = session() @@ -36,6 +42,17 @@ export const SessionView: Component = (props) => { } function getUserMessageText(messageId: string): string | null { + const normalizedMessage = messageStore().getMessage(messageId) + if (normalizedMessage && normalizedMessage.role === "user") { + const parts = normalizedMessage.partIds + .map((partId) => normalizedMessage.parts[partId]?.data) + .filter((part): part is ClientPart => Boolean(part)) + const textParts = parts.filter(isTextPart) + if (textParts.length > 0) { + return textParts.map((part) => part.text).join("\n") + } + } + const currentSession = session() if (!currentSession) return null @@ -45,7 +62,7 @@ export const SessionView: Component = (props) => { return null } - const textParts = targetMessage.parts.filter((p): p is ClientPart & { type: "text"; text: string } => p.type === "text") + const textParts = targetMessage.parts.filter(isTextPart) if (textParts.length === 0) { return null } @@ -129,12 +146,9 @@ export const SessionView: Component = (props) => { > {(s) => (
- props.toolCallId || props.toolCall?.id || "" - const pendingPermission = createMemo(() => props.toolCall.pendingPermission) + const store = createMemo(() => messageStoreBus.getOrCreate(props.instanceId)) + const permissionState = createMemo(() => store().getPermissionState(props.messageId, toolCallId() || props.toolCall?.id)) + const pendingPermission = createMemo(() => { + const state = permissionState() + if (state) { + return { permission: state.entry.permission, active: state.active } + } + return props.toolCall.pendingPermission + }) const expanded = () => (pendingPermission() ? true : isToolCallExpanded(toolCallId())) const toolOutputDefaultExpanded = createMemo(() => (preferences().toolOutputExpansion || "expanded") === "expanded") const diagnosticsDefaultExpanded = createMemo(() => (preferences().diagnosticsExpansion || "expanded") === "expanded") diff --git a/packages/ui/src/lib/global-cache.ts b/packages/ui/src/lib/global-cache.ts new file mode 100644 index 00000000..2c7b941b --- /dev/null +++ b/packages/ui/src/lib/global-cache.ts @@ -0,0 +1,97 @@ +interface CacheLocation { + instanceId?: string + sessionId?: string + scope?: string +} + +const GLOBAL_KEY = "GLOBAL" + +type CacheScope = Map +type ScopeCollection = Map +type SessionMap = Map +const cacheRoot = new Map() + +function resolveKey(value?: string) { + return value && value.length > 0 ? value : GLOBAL_KEY +} + +function resolveCacheScope(location: CacheLocation, createIfMissing: boolean): CacheScope | undefined { + const instanceKey = resolveKey(location.instanceId) + const sessionKey = resolveKey(location.sessionId) + const scopeKey = resolveKey(location.scope) + + let sessionMap = cacheRoot.get(instanceKey) + if (!sessionMap) { + if (!createIfMissing) return undefined + sessionMap = new Map() + cacheRoot.set(instanceKey, sessionMap) + } + + let scopeCollection = sessionMap.get(sessionKey) + if (!scopeCollection) { + if (!createIfMissing) return undefined + scopeCollection = new Map() + sessionMap.set(sessionKey, scopeCollection) + } + + let cacheScope = scopeCollection.get(scopeKey) + if (!cacheScope) { + if (!createIfMissing) return undefined + cacheScope = new Map() + scopeCollection.set(scopeKey, cacheScope) + } + + return cacheScope +} + +export function setGlobalCacheValue(location: CacheLocation, key: string, value: unknown): void { + const cacheScope = resolveCacheScope(location, true) + cacheScope?.set(key, value) +} + +export function getGlobalCacheValue(location: CacheLocation, key: string): T | undefined { + const cacheScope = resolveCacheScope(location, false) + return (cacheScope?.get(key) as T | undefined) ?? undefined +} + +export function deleteGlobalCacheValue(location: CacheLocation, key: string): void { + const cacheScope = resolveCacheScope(location, false) + cacheScope?.delete(key) +} + +export function clearGlobalCacheScope(location: CacheLocation): void { + const instanceKey = resolveKey(location.instanceId) + const sessionKey = resolveKey(location.sessionId) + const scopeKey = resolveKey(location.scope) + const sessionMap = cacheRoot.get(instanceKey) + if (!sessionMap) return + const scopeCollection = sessionMap.get(sessionKey) + if (!scopeCollection) return + scopeCollection.delete(scopeKey) + if (scopeCollection.size === 0) { + sessionMap.delete(sessionKey) + } + if (sessionMap.size === 0) { + cacheRoot.delete(instanceKey) + } +} + +export function clearGlobalCacheSession(instanceId?: string, sessionId?: string): void { + const instanceKey = resolveKey(instanceId) + const sessionKey = resolveKey(sessionId) + const sessionMap = cacheRoot.get(instanceKey) + if (!sessionMap) return + sessionMap.delete(sessionKey) + if (sessionMap.size === 0) { + cacheRoot.delete(instanceKey) + } +} + +export function clearGlobalCacheInstance(instanceId?: string): void { + const instanceKey = resolveKey(instanceId) + cacheRoot.delete(instanceKey) +} + +export function clearAllGlobalCache(): void { + cacheRoot.clear() +} diff --git a/packages/ui/src/lib/hooks/use-commands.ts b/packages/ui/src/lib/hooks/use-commands.ts index f413081c..1794af61 100644 --- a/packages/ui/src/lib/hooks/use-commands.ts +++ b/packages/ui/src/lib/hooks/use-commands.ts @@ -3,6 +3,7 @@ import type { Accessor } from "solid-js" import type { Preferences, ExpansionPreference } from "../../stores/preferences" import { createCommandRegistry, type Command } from "../commands" import { instances, activeInstanceId, setActiveInstanceId } from "../../stores/instances" +import type { ClientPart, MessageInfo } from "../../types/message" import { activeParentSessionId, activeSessionId as activeSessionMap, @@ -13,6 +14,8 @@ import { import { setSessionCompactionState } from "../../stores/session-compaction" import { showAlertDialog } from "../../stores/alerts" import type { Instance } from "../../types/instance" +import type { MessageRecord } from "../../stores/message-v2/types" +import { messageStoreBus } from "../../stores/message-v2/bus" export interface UseCommandsOptions { preferences: Accessor @@ -29,6 +32,18 @@ export interface UseCommandsOptions { getActiveSessionIdForInstance: () => string | null } +function extractUserTextFromRecord(record?: MessageRecord): string | null { + if (!record) return null + const parts = record.partIds + .map((partId) => record.parts[partId]?.data) + .filter((part): part is ClientPart => Boolean(part)) + const textParts = parts.filter((part): part is ClientPart & { type: "text"; text: string } => part.type === "text" && typeof (part as any).text === "string") + if (textParts.length === 0) { + return null + } + return textParts.map((part) => (part as any).text as string).join("\n") +} + export function useCommands(options: UseCommandsOptions) { const commandRegistry = createCommandRegistry() const [commands, setCommands] = createSignal([]) @@ -232,34 +247,56 @@ export function useCommands(options: UseCommandsOptions) { const session = sessions.find((s) => s.id === sessionId) if (!session) return - let after = 0 - const revert = session.revert + const store = messageStoreBus.getOrCreate(instance.id) + const messageIds = store.getSessionMessageIds(sessionId) + const infoMap = new Map() + messageIds.forEach((id) => { + const info = store.getMessageInfo(id) + if (info) infoMap.set(id, info) + }) - if (revert?.messageID) { - for (let i = session.messages.length - 1; i >= 0; i--) { - const msg = session.messages[i] - const info = session.messagesInfo.get(msg.id) - if (info?.id === revert.messageID) { - after = info.time?.created || 0 - break - } - } + const revertState = store.getSessionRevert(sessionId) ?? session.revert + let after = 0 + if (revertState?.messageID) { + const revertInfo = infoMap.get(revertState.messageID) ?? session.messagesInfo.get(revertState.messageID) + after = revertInfo?.time?.created || 0 } let messageID = "" - for (let i = session.messages.length - 1; i >= 0; i--) { - const msg = session.messages[i] - const info = session.messagesInfo.get(msg.id) - - if (msg.type === "user" && info?.time?.created) { + let restoredText: string | null = null + for (let i = messageIds.length - 1; i >= 0; i--) { + const id = messageIds[i] + const record = store.getMessage(id) + const info = infoMap.get(id) + if (record?.role === "user" && info?.time?.created) { if (after > 0 && info.time.created >= after) { continue } - messageID = msg.id + messageID = id + restoredText = extractUserTextFromRecord(record) break } } + if (!messageID) { + for (let i = session.messages.length - 1; i >= 0; i--) { + const msg = session.messages[i] + const info = session.messagesInfo.get(msg.id) + + if (msg.type === "user" && info?.time?.created) { + if (after > 0 && info.time.created >= after) { + continue + } + messageID = msg.id + const textParts = msg.parts.filter((p): p is ClientPart & { type: "text"; text: string } => p.type === "text" && typeof (p as any).text === "string") + if (textParts.length > 0) { + restoredText = textParts.map((p) => (p as any).text as string).join("\n") + } + break + } + } + } + if (!messageID) { showAlertDialog("Nothing to undo", { title: "No actions to undo", @@ -274,20 +311,27 @@ export function useCommands(options: UseCommandsOptions) { body: { messageID }, }) - const revertedMessage = session.messages.find((m) => m.id === messageID) - const revertedInfo = session.messagesInfo.get(messageID) - - if (revertedMessage && revertedInfo?.role === "user") { - const textParts = revertedMessage.parts.filter((p) => p.type === "text") - if (textParts.length > 0) { - const textarea = document.querySelector(".prompt-input") as HTMLTextAreaElement - if (textarea) { - textarea.value = textParts.map((p: any) => p.text).join("\n") - textarea.dispatchEvent(new Event("input", { bubbles: true })) - textarea.focus() + if (!restoredText) { + const revertedMessage = session.messages.find((m) => m.id === messageID) + const revertedInfo = session.messagesInfo.get(messageID) + if (revertedMessage && revertedInfo?.role === "user") { + const textParts = revertedMessage.parts.filter( + (p): p is ClientPart & { type: "text"; text: string } => p.type === "text" && typeof (p as any).text === "string", + ) + if (textParts.length > 0) { + restoredText = textParts.map((p) => (p as any).text as string).join("\n") } } } + + if (restoredText) { + const textarea = document.querySelector(".prompt-input") as HTMLTextAreaElement + if (textarea) { + textarea.value = restoredText + textarea.dispatchEvent(new Event("input", { bubbles: true })) + textarea.focus() + } + } } catch (error) { console.error("Failed to revert message:", error) showAlertDialog("Failed to revert message", { diff --git a/packages/ui/src/lib/scroll-cache.ts b/packages/ui/src/lib/scroll-cache.ts new file mode 100644 index 00000000..ff3f29c7 --- /dev/null +++ b/packages/ui/src/lib/scroll-cache.ts @@ -0,0 +1,53 @@ +import type { ScrollSnapshot } from "../stores/message-v2/types" + +interface ScrollCacheParams { + instanceId?: string + sessionId?: string + scope?: string +} + +const scrollCache = new Map() +const DEFAULT_SCOPE = "session" + +function resolve(value?: string) { + return value && value.length > 0 ? value : "GLOBAL" +} + +function makeKey(params: ScrollCacheParams) { + return `${resolve(params.instanceId)}:${resolve(params.sessionId)}:${params.scope ?? DEFAULT_SCOPE}` +} + +export function setScrollCache(params: ScrollCacheParams, snapshot: Omit) { + scrollCache.set(makeKey(params), { ...snapshot, updatedAt: Date.now() }) +} + +export function getScrollCache(params: ScrollCacheParams): ScrollSnapshot | undefined { + return scrollCache.get(makeKey(params)) +} + +export function clearScrollCacheScope(params: ScrollCacheParams) { + const key = makeKey(params) + scrollCache.delete(key) +} + +export function clearScrollCacheForSession(instanceId?: string, sessionId?: string) { + const match = `${resolve(instanceId)}:${resolve(sessionId)}:` + for (const key of scrollCache.keys()) { + if (key.startsWith(match)) { + scrollCache.delete(key) + } + } +} + +export function clearScrollCacheForInstance(instanceId?: string) { + const match = `${resolve(instanceId)}:` + for (const key of scrollCache.keys()) { + if (key.startsWith(match)) { + scrollCache.delete(key) + } + } +} + +export function clearAllScrollCache() { + scrollCache.clear() +} diff --git a/packages/ui/src/stores/message-v2/bridge.ts b/packages/ui/src/stores/message-v2/bridge.ts new file mode 100644 index 00000000..11def492 --- /dev/null +++ b/packages/ui/src/stores/message-v2/bridge.ts @@ -0,0 +1,175 @@ +import type { Permission } from "@opencode-ai/sdk" +import type { Message, MessageInfo, ClientPart } from "../../types/message" +import type { Session } from "../../types/session" +import { messageStoreBus } from "./bus" +import type { MessageStatus, SessionRevertState } from "./types" + +interface SessionMetadata { + id: string + title?: string + parentId?: string | null +} + +function resolveSessionMetadata(session?: Session | null): SessionMetadata | undefined { + if (!session) return undefined + return { + id: session.id, + title: session.title, + parentId: session.parentId ?? null, + } +} + +function normalizeStatus(status: Message["status"]): MessageStatus { + switch (status) { + case "sending": + case "sent": + case "streaming": + case "complete": + case "error": + return status + default: + return "complete" + } +} + +export function seedSessionMessagesV2( + instanceId: string, + session: Session | SessionMetadata, + messages: Message[], + messageInfos?: Map, +): void { + if (!session || !Array.isArray(messages)) return + const store = messageStoreBus.getOrCreate(instanceId) + const metadata: SessionMetadata = "id" in session ? { id: session.id, title: session.title, parentId: session.parentId ?? null } : session + const messageIds = messages.map((message) => message.id) + + store.addOrUpdateSession({ + id: metadata.id, + title: metadata.title, + parentId: metadata.parentId ?? null, + messageIds, + revert: (session as Session)?.revert ?? undefined, + }) + + messages.forEach((message) => { + store.upsertMessage({ + id: message.id, + sessionId: message.sessionId, + role: message.type, + status: normalizeStatus(message.status), + createdAt: message.timestamp, + updatedAt: message.timestamp, + parts: message.parts, + isEphemeral: message.status === "sending" || message.status === "streaming", + bumpRevision: false, + }) + const info = messageInfos?.get(message.id) + if (info) { + store.setMessageInfo(message.id, info) + } + }) + + if (messageInfos) { + store.rebuildUsage(metadata.id, messageInfos.values()) + } +} + +interface MessageInfoOptions { + status?: MessageStatus + bumpRevision?: boolean +} + +export function upsertMessageInfoV2(instanceId: string, info: MessageInfo | null | undefined, options?: MessageInfoOptions): void { + if (!info || typeof info.id !== "string" || typeof info.sessionID !== "string") { + return + } + const store = messageStoreBus.getOrCreate(instanceId) + const timeInfo = (info.time ?? {}) as { created?: number; completed?: number } + const createdAt = typeof timeInfo.created === "number" ? timeInfo.created : Date.now() + const completedAt = typeof timeInfo.completed === "number" ? timeInfo.completed : undefined + + store.upsertMessage({ + id: info.id, + sessionId: info.sessionID, + role: info.role === "user" ? "user" : "assistant", + status: options?.status ?? "complete", + createdAt, + updatedAt: completedAt ?? createdAt, + bumpRevision: Boolean(options?.bumpRevision), + }) + store.setMessageInfo(info.id, info) +} + +export function applyPartUpdateV2(instanceId: string, part: ClientPart | null | undefined): void { + if (!part || typeof part.messageID !== "string") { + return + } + const store = messageStoreBus.getOrCreate(instanceId) + store.applyPartUpdate({ + messageId: part.messageID, + part, + }) +} + +export function replaceMessageIdV2(instanceId: string, oldId: string, newId: string): void { + if (!oldId || !newId || oldId === newId) return + const store = messageStoreBus.getOrCreate(instanceId) + store.replaceMessageId({ oldId, newId }) +} + +function extractPermissionMessageId(permission: Permission): string | undefined { + return (permission as any).messageID || (permission as any).messageId +} + +function extractPermissionPartId(permission: Permission): string | undefined { + const metadata = (permission as any).metadata || {} + return ( + (permission as any).callID || + (permission as any).callId || + (permission as any).toolCallID || + (permission as any).toolCallId || + metadata.partId || + metadata.partID || + metadata.callID || + metadata.callId || + undefined + ) +} + +export function upsertPermissionV2(instanceId: string, permission: Permission): void { + if (!permission) return + const store = messageStoreBus.getOrCreate(instanceId) + store.upsertPermission({ + permission, + messageId: extractPermissionMessageId(permission), + partId: extractPermissionPartId(permission), + enqueuedAt: (permission as any).time?.created ?? Date.now(), + }) +} + +export function removePermissionV2(instanceId: string, permissionId: string): void { + if (!permissionId) return + const store = messageStoreBus.getOrCreate(instanceId) + store.removePermission(permissionId) +} + +export function ensureSessionMetadataV2(instanceId: string, session: Session | null | undefined): void { + if (!session) return + const store = messageStoreBus.getOrCreate(instanceId) + store.addOrUpdateSession({ + id: session.id, + title: session.title, + parentId: session.parentId ?? null, + messageIds: session.messages.map((message: Message) => message.id), + }) +} + +export function getSessionMetadataFromStore(session?: Session | null): SessionMetadata | undefined { + return resolveSessionMetadata(session ?? undefined) +} + +export function setSessionRevertV2(instanceId: string, sessionId: string, revert?: SessionRevertState | null): void { + if (!sessionId) return + const store = messageStoreBus.getOrCreate(instanceId) + store.setSessionRevert(sessionId, revert ?? null) +} diff --git a/packages/ui/src/stores/message-v2/bus.ts b/packages/ui/src/stores/message-v2/bus.ts new file mode 100644 index 00000000..ff027818 --- /dev/null +++ b/packages/ui/src/stores/message-v2/bus.ts @@ -0,0 +1,41 @@ +import { createInstanceMessageStore } from "./instance-store" +import type { InstanceMessageStore } from "./instance-store" + +class MessageStoreBus { + private stores = new Map() + + registerInstance(instanceId: string, store?: InstanceMessageStore): InstanceMessageStore { + if (this.stores.has(instanceId)) { + return this.stores.get(instanceId) as InstanceMessageStore + } + + const resolved = store ?? createInstanceMessageStore(instanceId) + this.stores.set(instanceId, resolved) + return resolved + } + + getInstance(instanceId: string): InstanceMessageStore | undefined { + return this.stores.get(instanceId) + } + + getOrCreate(instanceId: string): InstanceMessageStore { + return this.registerInstance(instanceId) + } + + unregisterInstance(instanceId: string) { + const store = this.stores.get(instanceId) + if (store) { + store.clearInstance() + } + this.stores.delete(instanceId) + } + + clearAll() { + for (const [instanceId, store] of this.stores.entries()) { + store.clearInstance() + this.stores.delete(instanceId) + } + } +} + +export const messageStoreBus = new MessageStoreBus() diff --git a/packages/ui/src/stores/message-v2/instance-store.ts b/packages/ui/src/stores/message-v2/instance-store.ts new file mode 100644 index 00000000..82cb5751 --- /dev/null +++ b/packages/ui/src/stores/message-v2/instance-store.ts @@ -0,0 +1,523 @@ +import { createStore, produce, reconcile } from "solid-js/store" +import type { SetStoreFunction } from "solid-js/store" +import type { ClientPart, MessageInfo } from "../../types/message" +import type { + InstanceMessageState, + MessageRecord, + MessageUpsertInput, + PartUpdateInput, + PendingPartEntry, + PermissionEntry, + ReplaceMessageIdOptions, + ScrollSnapshot, + SessionRecord, + SessionUpsertInput, + SessionUsageState, + UsageEntry, +} from "./types" + +function createInitialState(instanceId: string): InstanceMessageState { + return { + instanceId, + sessions: {}, + sessionOrder: [], + messages: {}, + messageInfo: {}, + pendingParts: {}, + permissions: { + queue: [], + active: null, + byMessage: {}, + }, + usage: {}, + scrollState: {}, + } +} + +function ensurePartId(messageId: string, part: ClientPart, index: number): string { + if (typeof part.id === "string" && part.id.length > 0) { + return part.id + } + return `${messageId}-part-${index}` +} + +function clonePart(part: ClientPart): ClientPart { + return JSON.parse(JSON.stringify(part)) as ClientPart +} + +function createEmptyUsageState(): SessionUsageState { + return { + entries: {}, + totalInputTokens: 0, + totalOutputTokens: 0, + totalReasoningTokens: 0, + totalCost: 0, + actualUsageTokens: 0, + latestMessageId: undefined, + } +} + +function extractUsageEntry(info: MessageInfo | undefined): UsageEntry | null { + if (!info || info.role !== "assistant") return null + const messageId = typeof info.id === "string" ? info.id : undefined + if (!messageId) return null + const tokens = info.tokens + if (!tokens) return null + const inputTokens = tokens.input ?? 0 + const outputTokens = tokens.output ?? 0 + const reasoningTokens = tokens.reasoning ?? 0 + const cacheReadTokens = tokens.cache?.read ?? 0 + const cacheWriteTokens = tokens.cache?.write ?? 0 + if (inputTokens === 0 && outputTokens === 0 && reasoningTokens === 0 && cacheReadTokens === 0 && cacheWriteTokens === 0) { + return null + } + const combinedTokens = info.summary ? outputTokens : inputTokens + cacheReadTokens + cacheWriteTokens + outputTokens + reasoningTokens + return { + messageId, + inputTokens, + outputTokens, + reasoningTokens, + cacheReadTokens, + cacheWriteTokens, + combinedTokens, + cost: info.cost ?? 0, + timestamp: info.time?.created ?? 0, + hasContextUsage: inputTokens + cacheReadTokens + cacheWriteTokens > 0, + } +} + +function applyUsageState(state: SessionUsageState, entry: UsageEntry | null) { + if (!entry) return + state.entries[entry.messageId] = entry + state.totalInputTokens += entry.inputTokens + state.totalOutputTokens += entry.outputTokens + state.totalReasoningTokens += entry.reasoningTokens + state.totalCost += entry.cost + if (!state.latestMessageId || entry.timestamp >= (state.entries[state.latestMessageId]?.timestamp ?? 0)) { + state.latestMessageId = entry.messageId + state.actualUsageTokens = entry.combinedTokens + } +} + +function removeUsageEntry(state: SessionUsageState, messageId: string | undefined) { + if (!messageId) return + const existing = state.entries[messageId] + if (!existing) return + state.totalInputTokens -= existing.inputTokens + state.totalOutputTokens -= existing.outputTokens + state.totalReasoningTokens -= existing.reasoningTokens + state.totalCost -= existing.cost + delete state.entries[messageId] + if (state.latestMessageId === messageId) { + state.latestMessageId = undefined + state.actualUsageTokens = 0 + let latest: UsageEntry | null = null + for (const candidate of Object.values(state.entries) as UsageEntry[]) { + if (!latest || candidate.timestamp >= latest.timestamp) { + latest = candidate + } + } + if (latest) { + state.latestMessageId = latest.messageId + state.actualUsageTokens = latest.combinedTokens + } + } +} + +function rebuildUsageStateFromInfos(infos: Iterable): SessionUsageState { + const usageState = createEmptyUsageState() + for (const info of infos) { + const entry = extractUsageEntry(info) + if (entry) { + applyUsageState(usageState, entry) + } + } + return usageState +} + +export interface InstanceMessageStore { + + instanceId: string + state: InstanceMessageState + setState: SetStoreFunction + addOrUpdateSession: (input: SessionUpsertInput) => void + upsertMessage: (input: MessageUpsertInput) => void + applyPartUpdate: (input: PartUpdateInput) => void + bufferPendingPart: (entry: PendingPartEntry) => void + flushPendingParts: (messageId: string) => void + replaceMessageId: (options: ReplaceMessageIdOptions) => void + setMessageInfo: (messageId: string, info: MessageInfo) => void + getMessageInfo: (messageId: string) => MessageInfo | undefined + upsertPermission: (entry: PermissionEntry) => void + removePermission: (permissionId: string) => void + getPermissionState: (messageId?: string, partId?: string) => { entry: PermissionEntry; active: boolean } | null + setSessionRevert: (sessionId: string, revert?: SessionRecord["revert"] | null) => void + getSessionRevert: (sessionId: string) => SessionRecord["revert"] | undefined | null + rebuildUsage: (sessionId: string, infos: Iterable) => void + getSessionUsage: (sessionId: string) => SessionUsageState | undefined + setScrollSnapshot: (sessionId: string, scope: string, snapshot: Omit) => void + getScrollSnapshot: (sessionId: string, scope: string) => ScrollSnapshot | undefined + getSessionMessageIds: (sessionId: string) => string[] + getMessage: (messageId: string) => MessageRecord | undefined + clearInstance: () => void +} + +export function createInstanceMessageStore(instanceId: string): InstanceMessageStore { + const [state, setState] = createStore(createInitialState(instanceId)) + + function withUsageState(sessionId: string, updater: (draft: SessionUsageState) => void) { + setState("usage", sessionId, (current) => { + const draft = current + ? { + ...current, + entries: { ...current.entries }, + } + : createEmptyUsageState() + updater(draft) + return draft + }) + } + + function updateUsageWithInfo(info: MessageInfo | undefined) { + if (!info || typeof info.sessionID !== "string") return + const messageId = typeof info.id === "string" ? info.id : undefined + if (!messageId) return + withUsageState(info.sessionID, (draft) => { + removeUsageEntry(draft, messageId) + const entry = extractUsageEntry(info) + if (entry) { + applyUsageState(draft, entry) + } + }) + } + + function rebuildUsage(sessionId: string, infos: Iterable) { + const usageState = rebuildUsageStateFromInfos(infos) + setState("usage", sessionId, usageState) + } + + function getSessionUsage(sessionId: string) { + return state.usage[sessionId] + } + + function ensureSessionEntry(sessionId: string): SessionRecord { + const existing = state.sessions[sessionId] + if (existing) { + return existing + } + + const now = Date.now() + const session: SessionRecord = { + id: sessionId, + createdAt: now, + updatedAt: now, + messageIds: [], + } + + setState("sessions", sessionId, session) + setState("sessionOrder", (order) => (order.includes(sessionId) ? order : [...order, sessionId])) + return session + } + + function addOrUpdateSession(input: SessionUpsertInput) { + const session = ensureSessionEntry(input.id) + const nextMessageIds = Array.isArray(input.messageIds) ? input.messageIds : session.messageIds + + setState("sessions", input.id, { + ...session, + title: input.title ?? session.title, + parentId: input.parentId ?? session.parentId ?? null, + updatedAt: Date.now(), + messageIds: nextMessageIds, + revert: input.revert ?? session.revert ?? null, + }) + } + + function insertMessageIntoSession(sessionId: string, messageId: string) { + ensureSessionEntry(sessionId) + setState("sessions", sessionId, "messageIds", (ids = []) => { + if (ids.includes(messageId)) { + return ids + } + return [...ids, messageId] + }) + } + + function normalizeParts(messageId: string, parts: ClientPart[] | undefined) { + if (!parts || parts.length === 0) { + return null + } + const map: MessageRecord["parts"] = {} + const ids: string[] = [] + + parts.forEach((part, index) => { + const id = ensurePartId(messageId, part, index) + const cloned = clonePart(part) + if (typeof cloned.version !== "number") { + cloned.version = 0 + } + map[id] = { + id, + data: cloned, + revision: cloned.version, + } + ids.push(id) + }) + + return { map, ids } + } + + function upsertMessage(input: MessageUpsertInput) { + const normalizedParts = normalizeParts(input.id, input.parts) + const shouldBump = Boolean(input.bumpRevision || normalizedParts) + const now = Date.now() + + setState("messages", input.id, (previous) => { + const revision = previous ? previous.revision + (shouldBump ? 1 : 0) : 0 + return { + id: input.id, + sessionId: input.sessionId, + role: input.role, + status: input.status, + createdAt: input.createdAt ?? previous?.createdAt ?? now, + updatedAt: input.updatedAt ?? now, + isEphemeral: input.isEphemeral ?? previous?.isEphemeral ?? false, + revision, + partIds: normalizedParts ? normalizedParts.ids : previous?.partIds ?? [], + parts: normalizedParts ? normalizedParts.map : previous?.parts ?? {}, + } + }) + + insertMessageIntoSession(input.sessionId, input.id) + flushPendingParts(input.id) + } + + function bufferPendingPart(entry: PendingPartEntry) { + setState("pendingParts", entry.messageId, (list = []) => [...list, entry]) + } + + function applyPartUpdate(input: PartUpdateInput) { + const message = state.messages[input.messageId] + if (!message) { + 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, + produce((draft: MessageRecord) => { + if (!draft.partIds.includes(partId)) { + draft.partIds = [...draft.partIds, partId] + } + const existing = draft.parts[partId] + const nextRevision = existing ? existing.revision + 1 : cloned.version ?? 0 + draft.parts[partId] = { + id: partId, + data: cloned, + revision: nextRevision, + } + draft.updatedAt = Date.now() + if (input.bumpRevision ?? true) { + draft.revision += 1 + } + }), + ) + } + + function flushPendingParts(messageId: string) { + const pending = state.pendingParts[messageId] + if (!pending || pending.length === 0) { + return + } + pending.forEach((entry) => applyPartUpdate({ messageId, part: entry.part })) + setState("pendingParts", (prev) => { + const next = { ...prev } + delete next[messageId] + return next + }) + } + + function replaceMessageId(options: ReplaceMessageIdOptions) { + if (options.oldId === options.newId) return + const existing = state.messages[options.oldId] + if (!existing) return + + const cloned: MessageRecord = { + ...existing, + id: options.newId, + isEphemeral: false, + updatedAt: Date.now(), + } + + setState("messages", options.newId, cloned) + setState("messages", (prev) => { + const next = { ...prev } + delete next[options.oldId] + return next + }) + + Object.values(state.sessions).forEach((session) => { + const index = session.messageIds.indexOf(options.oldId) + if (index === -1) return + setState("sessions", session.id, "messageIds", (ids) => { + const next = [...ids] + next[index] = options.newId + return next + }) + }) + + const infoEntry = state.messageInfo[options.oldId] + if (infoEntry) { + setState("messageInfo", options.newId, infoEntry) + setState("messageInfo", (prev) => { + const next = { ...prev } + delete next[options.oldId] + return next + }) + } + + const permissionMap = state.permissions.byMessage[options.oldId] + if (permissionMap) { + setState("permissions", "byMessage", options.newId, permissionMap) + setState("permissions", (prev) => { + const next = { ...prev } + const nextByMessage = { ...next.byMessage } + delete nextByMessage[options.oldId] + next.byMessage = nextByMessage + return next + }) + } + + const pending = state.pendingParts[options.oldId] + if (pending) { + setState("pendingParts", options.newId, pending) + setState("pendingParts", (prev) => { + const next = { ...prev } + delete next[options.oldId] + return next + }) + } + } + + function setMessageInfo(messageId: string, info: MessageInfo) { + if (!messageId) return + setState("messageInfo", messageId, info) + updateUsageWithInfo(info) + } + + function getMessageInfo(messageId: string) { + return state.messageInfo[messageId] + } + + function upsertPermission(entry: PermissionEntry) { + const messageKey = entry.messageId ?? "__global__" + const partKey = entry.partId ?? "__global__" + + setState( + "permissions", + produce((draft) => { + draft.byMessage[messageKey] = draft.byMessage[messageKey] ?? {} + draft.byMessage[messageKey][partKey] = entry + const existingIndex = draft.queue.findIndex((item) => item.permission.id === entry.permission.id) + if (existingIndex === -1) { + draft.queue.push(entry) + } else { + draft.queue[existingIndex] = entry + } + if (!draft.active || draft.active.permission.id === entry.permission.id) { + draft.active = entry + } + }), + ) + } + + function removePermission(permissionId: string) { + setState( + "permissions", + produce((draft) => { + draft.queue = draft.queue.filter((item) => item.permission.id !== permissionId) + if (draft.active?.permission.id === permissionId) { + draft.active = draft.queue[0] ?? null + } + Object.keys(draft.byMessage).forEach((messageKey) => { + const partEntries = draft.byMessage[messageKey] + Object.keys(partEntries).forEach((partKey) => { + if (partEntries[partKey].permission.id === permissionId) { + delete partEntries[partKey] + } + }) + if (Object.keys(partEntries).length === 0) { + delete draft.byMessage[messageKey] + } + }) + }), + ) + } + + function getPermissionState(messageId?: string, partId?: string) { + const messageKey = messageId ?? "__global__" + const partKey = partId ?? "__global__" + const entry = state.permissions.byMessage[messageKey]?.[partKey] + if (!entry) return null + const active = state.permissions.active?.permission.id === entry.permission.id + return { entry, active } + } + + function setSessionRevert(sessionId: string, revert?: SessionRecord["revert"] | null) { + if (!sessionId) return + ensureSessionEntry(sessionId) + setState("sessions", sessionId, "revert", revert ?? null) + } + + function getSessionRevert(sessionId: string) { + return state.sessions[sessionId]?.revert ?? null + } + + function makeScrollKey(sessionId: string, scope: string) { + return `${sessionId}:${scope}` + } + + function setScrollSnapshot(sessionId: string, scope: string, snapshot: Omit) { + const key = makeScrollKey(sessionId, scope) + setState("scrollState", key, { ...snapshot, updatedAt: Date.now() }) + } + + function getScrollSnapshot(sessionId: string, scope: string) { + const key = makeScrollKey(sessionId, scope) + return state.scrollState[key] + } + + function clearInstance() { + setState(reconcile(createInitialState(instanceId))) + } + + return { + instanceId, + state, + setState, + addOrUpdateSession, + upsertMessage, + applyPartUpdate, + bufferPendingPart, + flushPendingParts, + replaceMessageId, + setMessageInfo, + getMessageInfo, + upsertPermission, + removePermission, + getPermissionState, + setSessionRevert, + getSessionRevert, + rebuildUsage, + getSessionUsage, + setScrollSnapshot, + getScrollSnapshot, + getSessionMessageIds: (sessionId: string) => state.sessions[sessionId]?.messageIds ?? [], + getMessage: (messageId: string) => state.messages[messageId], + clearInstance, + } +} diff --git a/packages/ui/src/stores/message-v2/types.ts b/packages/ui/src/stores/message-v2/types.ts new file mode 100644 index 00000000..946c4dbf --- /dev/null +++ b/packages/ui/src/stores/message-v2/types.ts @@ -0,0 +1,138 @@ +import type { ClientPart, MessageInfo } from "../../types/message" +import type { Permission } from "@opencode-ai/sdk" + +export type MessageStatus = "sending" | "sent" | "streaming" | "complete" | "error" +export type MessageRole = "user" | "assistant" + +export interface NormalizedPartRecord { + id: string + data: ClientPart + revision: number +} + +export interface MessageRecord { + id: string + sessionId: string + role: MessageRole + status: MessageStatus + createdAt: number + updatedAt: number + revision: number + isEphemeral?: boolean + partIds: string[] + parts: Record +} + +export interface SessionRevertState { + messageID?: string + partID?: string + snapshot?: string + diff?: string +} + +export interface SessionRecord { + id: string + title?: string + parentId?: string | null + createdAt: number + updatedAt: number + messageIds: string[] + revert?: SessionRevertState | null +} + +export interface PendingPartEntry { + messageId: string + part: ClientPart + receivedAt: number +} + +export interface PermissionEntry { + permission: Permission + messageId?: string + partId?: string + enqueuedAt: number +} + +export interface InstancePermissionState { + queue: PermissionEntry[] + active: PermissionEntry | null + byMessage: Record> +} + +export interface ScrollSnapshot { + scrollTop: number + atBottom: boolean + updatedAt: number +} + +export interface UsageEntry { + messageId: string + inputTokens: number + outputTokens: number + reasoningTokens: number + cacheReadTokens: number + cacheWriteTokens: number + combinedTokens: number + cost: number + timestamp: number + hasContextUsage: boolean +} + +export interface SessionUsageState { + entries: Record + totalInputTokens: number + totalOutputTokens: number + totalReasoningTokens: number + totalCost: number + actualUsageTokens: number + latestMessageId?: string +} + +export interface InstanceMessageState { + instanceId: string + sessions: Record + sessionOrder: string[] + messages: Record + messageInfo: Record + pendingParts: Record + permissions: InstancePermissionState + usage: Record + scrollState: Record +} + +export interface SessionUpsertInput { + id: string + title?: string + parentId?: string | null + messageIds?: string[] + revert?: SessionRevertState | null +} + +export interface MessageUpsertInput { + id: string + sessionId: string + role: MessageRole + status: MessageStatus + parts?: ClientPart[] + createdAt?: number + updatedAt?: number + isEphemeral?: boolean + bumpRevision?: boolean +} + +export interface PartUpdateInput { + messageId: string + part: ClientPart + bumpRevision?: boolean +} + +export interface ReplaceMessageIdOptions { + oldId: string + newId: string +} + +export interface ScrollCacheKey { + instanceId: string + sessionId: string + scope: string +} diff --git a/packages/ui/src/stores/session-api.ts b/packages/ui/src/stores/session-api.ts index b9f28b09..ccc8d581 100644 --- a/packages/ui/src/stores/session-api.ts +++ b/packages/ui/src/stores/session-api.ts @@ -32,6 +32,7 @@ import { rebuildSessionUsage, updateSessionInfo, } from "./session-messages" +import { seedSessionMessagesV2 } from "./message-v2/bridge" interface SessionForkResponse { id: string @@ -610,6 +611,14 @@ async function loadMessages(instanceId: string, sessionId: string, force = false return next }) + const sessionForV2 = sessions().get(instanceId)?.get(sessionId) ?? { + id: sessionId, + title: session?.title, + parentId: session?.parentId ?? null, + revert: session?.revert, + } + seedSessionMessagesV2(instanceId, sessionForV2, messages, messagesInfo) + } catch (error) { console.error("Failed to load messages:", error) throw error diff --git a/packages/ui/src/stores/session-events.ts b/packages/ui/src/stores/session-events.ts index 8a466860..0f1e6378 100644 --- a/packages/ui/src/stores/session-events.ts +++ b/packages/ui/src/stores/session-events.ts @@ -34,6 +34,14 @@ import { } from "./session-messages" import { loadMessages } from "./session-api" import { setSessionCompactionState } from "./session-compaction" +import { + applyPartUpdateV2, + replaceMessageIdV2, + upsertMessageInfoV2, + upsertPermissionV2, + removePermissionV2, + setSessionRevertV2, +} from "./message-v2/bridge" interface TuiToastEvent { type: "tui.toast.show" @@ -195,6 +203,7 @@ function handleMessageUpdate(instanceId: string, event: MessageUpdateEvent | Mes index.partIndex.delete(oldId) index.partIndex.set(message.id, existingPartMap) } + replaceMessageIdV2(instanceId, oldId, message.id) } if (filteredSynthetics || replacedTemp) { @@ -212,6 +221,7 @@ function handleMessageUpdate(instanceId: string, event: MessageUpdateEvent | Mes /* mutations already applied above */ }) + applyPartUpdateV2(instanceId, part) updateSessionInfo(instanceId, part.sessionID) refreshPermissionsForSession(instanceId, part.sessionID) } else if (event.type === "message.updated") { @@ -270,6 +280,7 @@ function handleMessageUpdate(instanceId: string, event: MessageUpdateEvent | Mes index.partIndex.delete(oldId) index.partIndex.set(message.id, existingPartMap) } + replaceMessageIdV2(instanceId, oldId, message.id) } } else { const newMessage: any = { @@ -305,8 +316,11 @@ function handleMessageUpdate(instanceId: string, event: MessageUpdateEvent | Mes message.displayParts = computeDisplayParts(message, preferences().showThinkingBlocks) } + upsertMessageInfoV2(instanceId, info, { status: "complete" }) + session.messagesInfo.set(info.id, info) updateUsageFromMessageInfo(instanceId, info.sessionID, info) + withSession(instanceId, info.sessionID, () => { /* ensure reactivity */ }) @@ -359,6 +373,7 @@ function handleSessionUpdate(instanceId: string, event: EventSessionUpdated): vo next.set(instanceId, updated) return next }) + setSessionRevertV2(instanceId, info.id, info.revert ?? null) console.log(`[SSE] New session created: ${info.id}`, newSession) } else { @@ -391,6 +406,7 @@ function handleSessionUpdate(instanceId: string, event: EventSessionUpdated): vo next.set(instanceId, updated) return next }) + setSessionRevertV2(instanceId, info.id, info.revert ?? null) } } @@ -490,6 +506,7 @@ function handlePermissionUpdated(instanceId: string, event: EventPermissionUpdat console.log(`[SSE] Permission updated: ${permission.id} (${permission.type})`) addPermissionToQueue(instanceId, permission) + upsertPermissionV2(instanceId, permission) } function handlePermissionReplied(instanceId: string, event: EventPermissionReplied): void { @@ -498,6 +515,7 @@ function handlePermissionReplied(instanceId: string, event: EventPermissionRepli console.log(`[SSE] Permission replied: ${permissionID}`) removePermissionFromQueue(instanceId, permissionID) + removePermissionV2(instanceId, permissionID) } export { diff --git a/packages/ui/src/stores/session-status.ts b/packages/ui/src/stores/session-status.ts index b38cb936..54588104 100644 --- a/packages/ui/src/stores/session-status.ts +++ b/packages/ui/src/stores/session-status.ts @@ -1,7 +1,9 @@ import type { Session, SessionStatus } from "../types/session" import type { Message, MessageInfo } from "../types/message" +import type { MessageRecord } from "./message-v2/types" import { sessions } from "./sessions" import { isSessionCompactionActive } from "./session-compaction" +import { messageStoreBus } from "./message-v2/bus" function getSession(instanceId: string, sessionId: string): Session | null { const instanceSessions = sessions().get(instanceId) @@ -17,21 +19,49 @@ function isSessionCompacting(session: Session): boolean { return Boolean(compactingFlag) } -function getMessageTimestamp(session: Session, message?: Message): number { - if (!message) return Number.NEGATIVE_INFINITY - if (typeof message.timestamp === "number" && Number.isFinite(message.timestamp)) { - return message.timestamp +function getLatestInfoFromStore(instanceId: string, sessionId: string, role?: MessageInfo["role"]): MessageInfo | undefined { + const store = messageStoreBus.getOrCreate(instanceId) + const messageIds = store.getSessionMessageIds(sessionId) + let latest: MessageInfo | undefined + let latestTimestamp = Number.NEGATIVE_INFINITY + for (const id of messageIds) { + const info = store.getMessageInfo(id) + if (!info) continue + if (role && info.role !== role) continue + const timestamp = info.time?.created ?? 0 + if (timestamp >= latestTimestamp) { + latest = info + latestTimestamp = timestamp + } } - const info = session.messagesInfo.get(message.id) - return info?.time?.created ?? Number.NEGATIVE_INFINITY + return latest } -function getLastMessage(session: Session): Message | undefined { +function getLastMessageFromStore(instanceId: string, sessionId: string): MessageRecord | undefined { + const store = messageStoreBus.getOrCreate(instanceId) + const messageIds = store.getSessionMessageIds(sessionId) + let latest: MessageRecord | undefined + let latestTimestamp = Number.NEGATIVE_INFINITY + for (const id of messageIds) { + const record = store.getMessage(id) + if (!record) continue + const info = store.getMessageInfo(id) + const timestamp = info?.time?.created ?? record.createdAt ?? Number.NEGATIVE_INFINITY + if (timestamp >= latestTimestamp) { + latest = record + latestTimestamp = timestamp + } + } + return latest +} + +function getLegacyLastMessage(session: Session): Message | undefined { let latest: Message | undefined let latestTimestamp = Number.NEGATIVE_INFINITY for (const message of session.messages) { if (!message) continue - const timestamp = getMessageTimestamp(session, message) + const info = session.messagesInfo.get(message.id) + const timestamp = info?.time?.created ?? message.timestamp ?? Number.NEGATIVE_INFINITY if (timestamp >= latestTimestamp) { latest = message latestTimestamp = timestamp @@ -40,7 +70,7 @@ function getLastMessage(session: Session): Message | undefined { return latest } -function getLastMessageInfo(session: Session, role?: MessageInfo["role"]): MessageInfo | undefined { +function getLegacyLastMessageInfo(session: Session, role?: MessageInfo["role"]): MessageInfo | undefined { if (session.messagesInfo.size === 0) { return undefined } @@ -92,7 +122,28 @@ function isAssistantInfoPending(info?: MessageInfo): boolean { return completed < created } -function isAssistantStillGenerating(message: Message, info?: MessageInfo): boolean { +function isAssistantStillGeneratingRecord(record: MessageRecord, info?: MessageInfo): boolean { + if (record.role !== "assistant") { + return false + } + + if (record.status === "error") { + return false + } + + if (record.status === "streaming" || record.status === "sending") { + return true + } + + const completedAt = (info?.time as { completed?: number } | undefined)?.completed + if (completedAt !== undefined && completedAt !== null) { + return false + } + + return !(record.status === "complete" || record.status === "sent") +} + +function isAssistantStillGeneratingLegacy(message: Message, info?: MessageInfo): boolean { if (message.type !== "assistant") { return false } @@ -119,15 +170,20 @@ export function getSessionStatus(instanceId: string, sessionId: string): Session return "idle" } + const store = messageStoreBus.getOrCreate(instanceId) + if (isSessionCompactionActive(instanceId, sessionId) || isSessionCompacting(session)) { return "compacting" } - const latestUserInfo = getLastMessageInfo(session, "user") - const latestAssistantInfo = getLastMessageInfo(session, "assistant") - const lastMessage = getLastMessage(session) - if (!lastMessage) { - const latestInfo = getLastMessageInfo(session) + const latestUserInfo = getLatestInfoFromStore(instanceId, sessionId, "user") ?? getLegacyLastMessageInfo(session, "user") + const latestAssistantInfo = getLatestInfoFromStore(instanceId, sessionId, "assistant") ?? getLegacyLastMessageInfo(session, "assistant") + + const lastRecord = getLastMessageFromStore(instanceId, sessionId) + const legacyFallbackMessage = lastRecord ? undefined : getLegacyLastMessage(session) + + if (!lastRecord && !legacyFallbackMessage) { + const latestInfo = latestUserInfo ?? latestAssistantInfo ?? getLegacyLastMessageInfo(session) if (!latestInfo) { return "idle" } @@ -138,13 +194,22 @@ export function getSessionStatus(instanceId: string, sessionId: string): Session return infoCompleted ? "idle" : "working" } - if (lastMessage.type === "user") { - return "working" - } - - const infoForMessage = session.messagesInfo.get(lastMessage.id) ?? latestAssistantInfo - if (isAssistantStillGenerating(lastMessage, infoForMessage)) { - return "working" + if (lastRecord) { + if (lastRecord.role === "user") { + return "working" + } + const infoForRecord = store.getMessageInfo(lastRecord.id) ?? latestAssistantInfo + if (infoForRecord && isAssistantStillGeneratingRecord(lastRecord, infoForRecord)) { + return "working" + } + } else if (legacyFallbackMessage) { + if (legacyFallbackMessage.type === "user") { + return "working" + } + const infoForLegacy = session.messagesInfo.get(legacyFallbackMessage.id) ?? latestAssistantInfo + if (isAssistantStillGeneratingLegacy(legacyFallbackMessage, infoForLegacy)) { + return "working" + } } if (isAssistantInfoPending(latestAssistantInfo)) {