diff --git a/packages/ui/src/components/chat/ChatContainer.tsx b/packages/ui/src/components/chat/ChatContainer.tsx index 333c0504..e104485f 100644 --- a/packages/ui/src/components/chat/ChatContainer.tsx +++ b/packages/ui/src/components/chat/ChatContainer.tsx @@ -106,6 +106,7 @@ type ChatViewportProps = { isLoadingOlder: boolean; sessionIsWorking: boolean; streamingMessageId: string | null; + activeStreamingPhase: import('./message/types').StreamPhase | null; retryOverlay: { sessionId: string; message: string; @@ -135,6 +136,7 @@ const ChatViewport = React.memo(({ isLoadingOlder, sessionIsWorking, streamingMessageId, + activeStreamingPhase, retryOverlay, handleMessageContentChange, getAnimationHandlers, @@ -173,6 +175,7 @@ const ChatViewport = React.memo(({ messages={renderedMessages} sessionIsWorking={sessionIsWorking} activeStreamingMessageId={streamingMessageId} + activeStreamingPhase={activeStreamingPhase} retryOverlay={retryOverlay} onMessageContentChange={handleMessageContentChange} getAnimationHandlers={getAnimationHandlers} @@ -218,6 +221,7 @@ const ChatViewport = React.memo(({ && prev.isLoadingOlder === next.isLoadingOlder && prev.sessionIsWorking === next.sessionIsWorking && prev.streamingMessageId === next.streamingMessageId + && prev.activeStreamingPhase === next.activeStreamingPhase && prev.retryOverlay === next.retryOverlay && prev.handleMessageContentChange === next.handleMessageContentChange && prev.getAnimationHandlers === next.getAnimationHandlers @@ -296,6 +300,15 @@ export const ChatContainer: React.FC = () => { [currentSessionId], ), ); + const activeStreamingPhase = useStreamingStore( + React.useCallback( + (s) => { + if (!streamingMessageId) return null; + return s.messageStreamStates.get(streamingMessageId)?.phase ?? null; + }, + [streamingMessageId], + ), + ); const sessionMessageCount = useSessionMessageCount(currentSessionId ?? ''); const [suspendDetachedTailUpdates, setSuspendDetachedTailUpdates] = React.useState(false); // Messages from sync system @@ -803,6 +816,7 @@ export const ChatContainer: React.FC = () => { isLoadingOlder={timelineController.isLoadingOlder} sessionIsWorking={sessionIsWorking} streamingMessageId={streamingMessageId} + activeStreamingPhase={activeStreamingPhase} retryOverlay={retryOverlay} handleMessageContentChange={handleMessageContentChange} getAnimationHandlers={getAnimationHandlers} diff --git a/packages/ui/src/components/chat/ChatMessage.tsx b/packages/ui/src/components/chat/ChatMessage.tsx index ec76418b..68678170 100644 --- a/packages/ui/src/components/chat/ChatMessage.tsx +++ b/packages/ui/src/components/chat/ChatMessage.tsx @@ -8,7 +8,6 @@ import { useConfigStore } from '@/stores/useConfigStore'; import { useFeatureFlagsStore } from '@/stores/useFeatureFlagsStore'; import { useUIStore } from '@/stores/useUIStore'; import { useContextStore } from '@/stores/contextStore'; -import { useStreamingStore } from '@/sync/streaming'; import { useSessionUIStore } from '@/sync/session-ui-store'; import { useSelectionStore } from '@/sync/selection-store'; import * as sessionActions from '@/sync/session-actions'; @@ -131,6 +130,7 @@ interface ChatMessageProps { turnGroupingContext?: TurnGroupingContext; assistantHeaderMessageId?: string; isInActiveTurn?: boolean; + activeStreamingPhase?: StreamPhase | null; animateUserOnMount?: boolean; onUserAnimationConsumed?: (messageId: string) => void; } @@ -144,6 +144,7 @@ const ChatMessage: React.FC = ({ turnGroupingContext, assistantHeaderMessageId, isInActiveTurn = false, + activeStreamingPhase = null, animateUserOnMount = false, onUserAnimationConsumed, }) => { @@ -152,13 +153,6 @@ const ChatMessage: React.FC = ({ const messageContainerRef = React.useRef(null); const currentSessionId = useSessionUIStore((s) => s.currentSessionId); - const streamState = useStreamingStore((s) => s.messageStreamStates.get(message.info.id)); - const lifecyclePhase = isInActiveTurn ? (streamState?.phase ?? null) : null; - - const msgSessionId = (message.info as { sessionID?: string }).sessionID ?? currentSessionId ?? null; - const streamingMsgForSession = useStreamingStore((s) => msgSessionId ? s.streamingMessageIds.get(msgSessionId) ?? null : null); - const isStreamingMessage = isInActiveTurn ? streamingMsgForSession === message.info.id : false; - const hasActiveStreamInSession = typeof streamingMsgForSession === 'string' && streamingMsgForSession.length > 0; const getAgentModelForSession = useSelectionStore((s) => s.getAgentModelForSession); const getSessionModelSelection = useSelectionStore((s) => s.getSessionModelSelection); @@ -166,13 +160,8 @@ const ChatMessage: React.FC = ({ const forkFromMessage = sessionActions.forkFromMessage; streamPerfCount('ui.chat_message.render'); - if (isStreamingMessage) { + if (isInActiveTurn) { streamPerfCount('ui.chat_message.render.streaming'); - } else if (hasActiveStreamInSession) { - streamPerfCount('ui.chat_message.render.static_during_stream'); - if (!isInActiveTurn) { - streamPerfCount('ui.chat_message.render.static_outside_active_turn_during_stream'); - } } const providers = useConfigStore.getState().providers; @@ -596,11 +585,11 @@ const ChatMessage: React.FC = ({ if (isMessageCompleted) { return 'completed'; } - if (lifecyclePhase) { - return lifecyclePhase; + if (isInActiveTurn) { + return activeStreamingPhase ?? 'streaming'; } - return isStreamingMessage ? 'streaming' : 'completed'; - }, [isMessageCompleted, lifecyclePhase, isStreamingMessage]); + return 'completed'; + }, [activeStreamingPhase, isInActiveTurn, isMessageCompleted]); React.useEffect(() => { if (!isUser || !animateUserOnMount) { @@ -1148,6 +1137,7 @@ export default React.memo(ChatMessage, (prev, next) => { && prev.turnGroupingContext === next.turnGroupingContext && prev.assistantHeaderMessageId === next.assistantHeaderMessageId && prev.isInActiveTurn === next.isInActiveTurn + && prev.activeStreamingPhase === next.activeStreamingPhase && prev.animateUserOnMount === next.animateUserOnMount && prev.onUserAnimationConsumed === next.onUserAnimationConsumed; }); diff --git a/packages/ui/src/components/chat/MessageList.tsx b/packages/ui/src/components/chat/MessageList.tsx index 2394219e..50a08020 100644 --- a/packages/ui/src/components/chat/MessageList.tsx +++ b/packages/ui/src/components/chat/MessageList.tsx @@ -14,6 +14,7 @@ import { useUIStore } from '@/stores/useUIStore'; import { FadeInDisabledProvider } from './message/FadeInOnReveal'; import { hasPendingUserSendAnimation, consumePendingUserSendAnimation } from '@/lib/userSendAnimation'; import { streamPerfCount, streamPerfMeasure } from '@/stores/utils/streamDebug'; +import type { StreamPhase } from './message/types'; const MESSAGE_LIST_VIRTUALIZE_THRESHOLD = 40; const MESSAGE_LIST_OVERSCAN = 6; @@ -304,6 +305,7 @@ interface MessageListProps { messages: ChatMessageEntry[]; sessionIsWorking?: boolean; activeStreamingMessageId?: string | null; + activeStreamingPhase?: StreamPhase | null; retryOverlay?: { sessionId: string; message: string; @@ -347,6 +349,7 @@ interface MessageRowProps { turnGroupingContext?: TurnGroupingContext; assistantHeaderMessageId?: string; isInActiveTurn?: boolean; + activeStreamingPhase?: StreamPhase | null; animateUserOnMount?: boolean; onUserAnimationConsumed?: (messageId: string) => void; onContentChange: (reason?: ContentChangeReason) => void; @@ -361,6 +364,7 @@ const MessageRow = React.memo(({ turnGroupingContext, assistantHeaderMessageId, isInActiveTurn, + activeStreamingPhase, animateUserOnMount, onUserAnimationConsumed, onContentChange, @@ -380,6 +384,7 @@ const MessageRow = React.memo(({ turnGroupingContext={turnGroupingContext} assistantHeaderMessageId={assistantHeaderMessageId} isInActiveTurn={isInActiveTurn} + activeStreamingPhase={activeStreamingPhase} /> ); }, (prev, next) => { @@ -404,6 +409,7 @@ const MessageRow = React.memo(({ && prevTurn?.activityParts === nextTurn?.activityParts && prev.assistantHeaderMessageId === next.assistantHeaderMessageId && prev.isInActiveTurn === next.isInActiveTurn + && prev.activeStreamingPhase === next.activeStreamingPhase && prev.animationHandlers?.onChunk === next.animationHandlers?.onChunk && prev.animationHandlers?.onComplete === next.animationHandlers?.onComplete && prev.animationHandlers?.onStreamingCandidate === next.animationHandlers?.onStreamingCandidate @@ -430,6 +436,7 @@ interface TurnBlockProps { shouldAnimateUserMessage: (message: ChatMessageEntry) => boolean; onUserAnimationConsumed: (messageId: string) => void; activeStreamingMessageId?: string | null; + activeStreamingPhase?: StreamPhase | null; } const TurnBlock: React.FC = ({ @@ -447,6 +454,7 @@ const TurnBlock: React.FC = ({ shouldAnimateUserMessage, onUserAnimationConsumed, activeStreamingMessageId, + activeStreamingPhase, }) => { const turnUiState = turnUiStates.get(turn.turnId) ?? { isExpanded: defaultActivityExpanded }; const handleToggleTurnGroup = React.useCallback(() => { @@ -656,6 +664,7 @@ const TurnBlock: React.FC = ({ turnGroupingContext={turnGroupingContext} assistantHeaderMessageId={assistantHeaderMessageId} isInActiveTurn={Boolean(streamingAssistantMessageId) && message.info.id === streamingAssistantMessageId} + activeStreamingPhase={message.info.id === streamingAssistantMessageId ? activeStreamingPhase : null} animateUserOnMount={shouldAnimateUserMessage(message)} onUserAnimationConsumed={onUserAnimationConsumed} onContentChange={onMessageContentChange} @@ -681,6 +690,7 @@ const TurnBlock: React.FC = ({ turnUiState.isExpanded, turnGroupingContextBase, streamingAssistantMessageId, + activeStreamingPhase, visibleAssistantMessages, visibleAssistantIds, activityOwnerMessageId, @@ -717,6 +727,7 @@ interface UngroupedMessageRowProps { shouldAnimateUserMessage: (message: ChatMessageEntry) => boolean; onUserAnimationConsumed: (messageId: string) => void; activeStreamingMessageId?: string | null; + activeStreamingPhase?: StreamPhase | null; } const UngroupedMessageRow: React.FC = React.memo(({ @@ -729,6 +740,7 @@ const UngroupedMessageRow: React.FC = React.memo(({ shouldAnimateUserMessage, onUserAnimationConsumed, activeStreamingMessageId, + activeStreamingPhase, }) => { return ( = React.memo(({ animationHandlers={getAnimationHandlers(message.info.id)} scrollToBottom={scrollToBottom} isInActiveTurn={Boolean(activeStreamingMessageId) && message.info.id === activeStreamingMessageId} + activeStreamingPhase={message.info.id === activeStreamingMessageId ? activeStreamingPhase : null} /> ); }, (prev, next) => { @@ -752,7 +765,8 @@ const UngroupedMessageRow: React.FC = React.memo(({ && prev.scrollToBottom === next.scrollToBottom && prev.shouldAnimateUserMessage === next.shouldAnimateUserMessage && prev.onUserAnimationConsumed === next.onUserAnimationConsumed - && prev.activeStreamingMessageId === next.activeStreamingMessageId; + && prev.activeStreamingMessageId === next.activeStreamingMessageId + && prev.activeStreamingPhase === next.activeStreamingPhase; }); UngroupedMessageRow.displayName = 'UngroupedMessageRow'; @@ -771,6 +785,7 @@ interface MessageListEntryProps { shouldAnimateUserMessage: (message: ChatMessageEntry) => boolean; onUserAnimationConsumed: (messageId: string) => void; activeStreamingMessageId?: string | null; + activeStreamingPhase?: StreamPhase | null; } const turnContainsMessageId = (turn: TurnRecord, messageId: string | null | undefined): boolean => { @@ -799,6 +814,7 @@ const MessageListEntry: React.FC = React.memo(({ shouldAnimateUserMessage, onUserAnimationConsumed, activeStreamingMessageId, + activeStreamingPhase, }) => { if (entry.kind === 'ungrouped') { return ( @@ -812,6 +828,7 @@ const MessageListEntry: React.FC = React.memo(({ shouldAnimateUserMessage={shouldAnimateUserMessage} onUserAnimationConsumed={onUserAnimationConsumed} activeStreamingMessageId={activeStreamingMessageId} + activeStreamingPhase={activeStreamingPhase} /> ); } @@ -828,6 +845,7 @@ const MessageListEntry: React.FC = React.memo(({ shouldAnimateUserMessage={shouldAnimateUserMessage} onUserAnimationConsumed={onUserAnimationConsumed} activeStreamingMessageId={activeStreamingMessageId} + activeStreamingPhase={activeStreamingPhase} onMessageContentChange={onMessageContentChange} getAnimationHandlers={getAnimationHandlers} scrollToBottom={scrollToBottom} @@ -871,6 +889,22 @@ function areMessageListEntryPropsEqual(prevProps: MessageListEntryProps, nextPro } } + if (prevProps.activeStreamingPhase !== nextProps.activeStreamingPhase) { + const prevAffected = turnContainsMessageId(prevEntry.turn, prevProps.activeStreamingMessageId); + const nextAffected = turnContainsMessageId(nextEntry.turn, nextProps.activeStreamingMessageId); + if (prevAffected || nextAffected) { + return false; + } + } + + if (prevProps.activeStreamingPhase !== nextProps.activeStreamingPhase) { + const prevAffected = turnContainsMessageId(prevEntry.turn, prevProps.activeStreamingMessageId); + const nextAffected = turnContainsMessageId(nextEntry.turn, nextProps.activeStreamingMessageId); + if (prevAffected || nextAffected) { + return false; + } + } + return true; } @@ -884,6 +918,20 @@ function areMessageListEntryPropsEqual(prevProps: MessageListEntryProps, nextPro } } + if (prevProps.activeStreamingPhase !== nextProps.activeStreamingPhase) { + const messageId = prevEntry.message.info.id; + if (prevProps.activeStreamingMessageId === messageId || nextProps.activeStreamingMessageId === messageId) { + return false; + } + } + + if (prevProps.activeStreamingPhase !== nextProps.activeStreamingPhase) { + const messageId = prevEntry.message.info.id; + if (prevProps.activeStreamingMessageId === messageId || nextProps.activeStreamingMessageId === messageId) { + return false; + } + } + return ( areRenderRelevantMessagesEqual(prevEntry.message, nextEntry.message) && areOptionalRenderRelevantMessagesEqual(prevEntry.previousMessage, nextEntry.previousMessage) @@ -912,7 +960,8 @@ const StaticHistoryList: React.FC<{ chatRenderMode: 'sorted' | 'live'; shouldAnimateUserMessage: (message: ChatMessageEntry) => boolean; onUserAnimationConsumed: (messageId: string) => void; -}> = React.memo(({ entries, shouldVirtualize, virtualRows, totalSize, measureElement, contentRef, onMessageContentChange, getAnimationHandlers, scrollToBottom, stickyUserHeader, defaultActivityExpanded, turnUiStates, onToggleTurnGroup, chatRenderMode, shouldAnimateUserMessage, onUserAnimationConsumed }) => { + activeStreamingPhase?: StreamPhase | null; +}> = React.memo(({ entries, shouldVirtualize, virtualRows, totalSize, measureElement, contentRef, onMessageContentChange, getAnimationHandlers, scrollToBottom, stickyUserHeader, defaultActivityExpanded, turnUiStates, onToggleTurnGroup, chatRenderMode, shouldAnimateUserMessage, onUserAnimationConsumed, activeStreamingPhase }) => { const renderEntry = React.useCallback((entry: RenderEntry) => { return ( ); - }, [chatRenderMode, defaultActivityExpanded, getAnimationHandlers, onMessageContentChange, onToggleTurnGroup, onUserAnimationConsumed, scrollToBottom, shouldAnimateUserMessage, stickyUserHeader, turnUiStates]); + }, [activeStreamingPhase, chatRenderMode, defaultActivityExpanded, getAnimationHandlers, onMessageContentChange, onToggleTurnGroup, onUserAnimationConsumed, scrollToBottom, shouldAnimateUserMessage, stickyUserHeader, turnUiStates]); const paddingTop = shouldVirtualize && virtualRows.length > 0 ? virtualRows[0]?.start ?? 0 @@ -995,7 +1045,8 @@ const StaticHistoryList: React.FC<{ && prevProps.onToggleTurnGroup === nextProps.onToggleTurnGroup && prevProps.chatRenderMode === nextProps.chatRenderMode && prevProps.shouldAnimateUserMessage === nextProps.shouldAnimateUserMessage - && prevProps.onUserAnimationConsumed === nextProps.onUserAnimationConsumed; + && prevProps.onUserAnimationConsumed === nextProps.onUserAnimationConsumed + && prevProps.activeStreamingPhase === nextProps.activeStreamingPhase; }); StaticHistoryList.displayName = 'StaticHistoryList'; @@ -1014,6 +1065,7 @@ const StreamingTailContent: React.FC<{ shouldAnimateUserMessage: (message: ChatMessageEntry) => boolean; onUserAnimationConsumed: (messageId: string) => void; activeStreamingMessageId?: string | null; + activeStreamingPhase?: StreamPhase | null; }> = React.memo(({ entry, onMessageContentChange, @@ -1028,6 +1080,7 @@ const StreamingTailContent: React.FC<{ shouldAnimateUserMessage, onUserAnimationConsumed, activeStreamingMessageId, + activeStreamingPhase, }) => { return ( ); }, (prev, next) => { @@ -1059,7 +1113,8 @@ const StreamingTailContent: React.FC<{ && prev.chatRenderMode === next.chatRenderMode && prev.shouldAnimateUserMessage === next.shouldAnimateUserMessage && prev.onUserAnimationConsumed === next.onUserAnimationConsumed - && prev.activeStreamingMessageId === next.activeStreamingMessageId; + && prev.activeStreamingMessageId === next.activeStreamingMessageId + && prev.activeStreamingPhase === next.activeStreamingPhase; }); StreamingTailContent.displayName = 'StreamingTailContent'; @@ -1071,6 +1126,7 @@ const MessageList = React.forwardRef(({ messages, sessionIsWorking = false, activeStreamingMessageId = null, + activeStreamingPhase = null, retryOverlay = null, onMessageContentChange, getAnimationHandlers, @@ -1760,6 +1816,7 @@ const MessageList = React.forwardRef(({ chatRenderMode={chatRenderMode} shouldAnimateUserMessage={shouldAnimateUserMessage} onUserAnimationConsumed={onUserAnimationConsumed} + activeStreamingPhase={activeStreamingPhase} /> {trailingStreamingEntry ? ( (({ shouldAnimateUserMessage={shouldAnimateUserMessage} onUserAnimationConsumed={onUserAnimationConsumed} activeStreamingMessageId={activeStreamingMessageId} + activeStreamingPhase={activeStreamingPhase} /> ) : null} diff --git a/packages/ui/src/sync/sync-context.tsx b/packages/ui/src/sync/sync-context.tsx index 035531b8..5dc0846b 100644 --- a/packages/ui/src/sync/sync-context.tsx +++ b/packages/ui/src/sync/sync-context.tsx @@ -36,7 +36,14 @@ type SyncSystem = { directory: string } -const SyncContext = createContext(null) +const SYNC_CONTEXT_GLOBAL_KEY = "__openchamber_sync_context__" +type SyncGlobal = typeof globalThis & { + [SYNC_CONTEXT_GLOBAL_KEY]?: React.Context +} + +const syncGlobal = globalThis as SyncGlobal +const SyncContext = syncGlobal[SYNC_CONTEXT_GLOBAL_KEY] ?? createContext(null) +syncGlobal[SYNC_CONTEXT_GLOBAL_KEY] = SyncContext function useSyncSystem() { const ctx = useContext(SyncContext) @@ -1072,9 +1079,3 @@ const EMPTY_MESSAGES: Message[] = [] const EMPTY_PARTS: Part[] = [] const EMPTY_PERMISSION_REQUESTS: PermissionRequest[] = [] const EMPTY_QUESTION_REQUESTS: QuestionRequest[] = [] - -if (import.meta.hot) { - import.meta.hot.accept(() => { - import.meta.hot?.invalidate() - }) -}