diff --git a/autogpt_platform/frontend/src/app/(platform)/copilot/components/CopilotShell/CopilotShell.tsx b/autogpt_platform/frontend/src/app/(platform)/copilot/components/CopilotShell/CopilotShell.tsx index fb22640302..8c9f9d528c 100644 --- a/autogpt_platform/frontend/src/app/(platform)/copilot/components/CopilotShell/CopilotShell.tsx +++ b/autogpt_platform/frontend/src/app/(platform)/copilot/components/CopilotShell/CopilotShell.tsx @@ -3,7 +3,7 @@ import { ChatLoader } from "@/components/contextual/Chat/components/ChatLoader/ChatLoader"; import { NAVBAR_HEIGHT_PX } from "@/lib/constants"; import type { ReactNode } from "react"; -import { useEffect } from "react"; +import { useCallback, useEffect } from "react"; import { useCopilotStore } from "../../copilot-page-store"; import { DesktopSidebar } from "./components/DesktopSidebar/DesktopSidebar"; import { LoadingState } from "./components/LoadingState/LoadingState"; @@ -25,10 +25,12 @@ export function CopilotShell({ children }: Props) { sessions, currentSessionId, handleSelectSession, + performSelectSession, handleOpenDrawer, handleCloseDrawer, handleDrawerOpenChange, handleNewChat, + performNewChat, hasNextPage, isFetchingNextPage, fetchNextPage, @@ -36,22 +38,71 @@ export function CopilotShell({ children }: Props) { } = useCopilotShell(); const setNewChatHandler = useCopilotStore((s) => s.setNewChatHandler); + const setNewChatWithInterruptHandler = useCopilotStore( + (s) => s.setNewChatWithInterruptHandler, + ); + const setSelectSessionHandler = useCopilotStore( + (s) => s.setSelectSessionHandler, + ); + const setSelectSessionWithInterruptHandler = useCopilotStore( + (s) => s.setSelectSessionWithInterruptHandler, + ); const requestNewChat = useCopilotStore((s) => s.requestNewChat); + const requestSelectSession = useCopilotStore((s) => s.requestSelectSession); + + const stableHandleNewChat = useCallback(handleNewChat, [handleNewChat]); + const stablePerformNewChat = useCallback(performNewChat, [performNewChat]); useEffect( - function registerNewChatHandler() { - setNewChatHandler(handleNewChat); + function registerNewChatHandlers() { + setNewChatHandler(stableHandleNewChat); + setNewChatWithInterruptHandler(stablePerformNewChat); return function cleanup() { setNewChatHandler(null); + setNewChatWithInterruptHandler(null); }; }, - [handleNewChat], + [ + stableHandleNewChat, + stablePerformNewChat, + setNewChatHandler, + setNewChatWithInterruptHandler, + ], + ); + + const stableHandleSelectSession = useCallback(handleSelectSession, [ + handleSelectSession, + ]); + + const stablePerformSelectSession = useCallback(performSelectSession, [ + performSelectSession, + ]); + + useEffect( + function registerSelectSessionHandlers() { + setSelectSessionHandler(stableHandleSelectSession); + setSelectSessionWithInterruptHandler(stablePerformSelectSession); + return function cleanup() { + setSelectSessionHandler(null); + setSelectSessionWithInterruptHandler(null); + }; + }, + [ + stableHandleSelectSession, + stablePerformSelectSession, + setSelectSessionHandler, + setSelectSessionWithInterruptHandler, + ], ); function handleNewChatClick() { requestNewChat(); } + function handleSessionClick(sessionId: string) { + requestSelectSession(sessionId); + } + if (!isLoggedIn) { return (
@@ -72,7 +123,7 @@ export function CopilotShell({ children }: Props) { isLoading={isLoading} hasNextPage={hasNextPage} isFetchingNextPage={isFetchingNextPage} - onSelectSession={handleSelectSession} + onSelectSession={handleSessionClick} onFetchNextPage={fetchNextPage} onNewChat={handleNewChatClick} hasActiveSession={Boolean(hasActiveSession)} @@ -94,7 +145,7 @@ export function CopilotShell({ children }: Props) { isLoading={isLoading} hasNextPage={hasNextPage} isFetchingNextPage={isFetchingNextPage} - onSelectSession={handleSelectSession} + onSelectSession={handleSessionClick} onFetchNextPage={fetchNextPage} onNewChat={handleNewChatClick} onClose={handleCloseDrawer} diff --git a/autogpt_platform/frontend/src/app/(platform)/copilot/components/CopilotShell/useCopilotShell.ts b/autogpt_platform/frontend/src/app/(platform)/copilot/components/CopilotShell/useCopilotShell.ts index a3aa0b55b2..3154df2975 100644 --- a/autogpt_platform/frontend/src/app/(platform)/copilot/components/CopilotShell/useCopilotShell.ts +++ b/autogpt_platform/frontend/src/app/(platform)/copilot/components/CopilotShell/useCopilotShell.ts @@ -1,17 +1,20 @@ "use client"; import { + getGetV2GetSessionQueryKey, getGetV2ListSessionsQueryKey, useGetV2GetSession, } from "@/app/api/__generated__/endpoints/chat/chat"; import type { SessionSummaryResponse } from "@/app/api/__generated__/models/sessionSummaryResponse"; import { okData } from "@/app/api/helpers"; +import { useChatStore } from "@/components/contextual/Chat/chat-store"; import { useBreakpoint } from "@/lib/hooks/useBreakpoint"; import { useSupabase } from "@/lib/supabase/hooks/useSupabase"; import { useQueryClient } from "@tanstack/react-query"; import { parseAsString, useQueryState } from "nuqs"; import { usePathname, useSearchParams } from "next/navigation"; import { useEffect, useRef, useState } from "react"; +import { useCopilotStore } from "../../copilot-page-store"; import { useMobileDrawer } from "./components/MobileDrawer/useMobileDrawer"; import { useSessionsPagination } from "./components/SessionsList/useSessionsPagination"; import { @@ -73,6 +76,19 @@ export function useCopilotShell() { Map >(new Map()); + const [optimisticSessionId, setOptimisticSessionId] = useState( + null, + ); + + useEffect( + function clearOptimisticWhenUrlMatches() { + if (optimisticSessionId && currentSessionId === optimisticSessionId) { + setOptimisticSessionId(null); + } + }, + [currentSessionId, optimisticSessionId], + ); + // Mark as auto-selected when sessionId is in URL useEffect(() => { if (paramSessionId && !hasAutoSelectedRef.current) { @@ -142,7 +158,9 @@ export function useCopilotShell() { const visibleSessions = filterVisibleSessions(sessions); const sidebarSelectedSessionId = - isOnHomepage && !paramSessionId ? null : currentSessionId; + isOnHomepage && !paramSessionId && !optimisticSessionId + ? null + : optimisticSessionId || currentSessionId; const isReadyToShowContent = isOnHomepage ? true @@ -155,8 +173,89 @@ export function useCopilotShell() { hasAutoSelectedSession, ); - function handleSelectSession(sessionId: string) { + const stopStream = useChatStore((s) => s.stopStream); + const onStreamComplete = useChatStore((s) => s.onStreamComplete); + const setIsSwitchingSession = useCopilotStore((s) => s.setIsSwitchingSession); + + async function performSelectSession(sessionId: string) { + if (sessionId === currentSessionId) return; + + const sourceSessionId = currentSessionId; + + if (sourceSessionId) { + setIsSwitchingSession(true); + + await new Promise(function waitForStreamComplete(resolve) { + const unsubscribe = onStreamComplete( + function handleComplete(completedId) { + if (completedId === sourceSessionId) { + clearTimeout(timeout); + unsubscribe(); + resolve(); + } + }, + ); + const timeout = setTimeout(function handleTimeout() { + unsubscribe(); + resolve(); + }, 3000); + stopStream(sourceSessionId); + }); + + queryClient.invalidateQueries({ + queryKey: getGetV2GetSessionQueryKey(sourceSessionId), + }); + } + + setOptimisticSessionId(sessionId); setUrlSessionId(sessionId, { shallow: false }); + setIsSwitchingSession(false); + if (isMobile) handleCloseDrawer(); + } + + function handleSelectSession(sessionId: string) { + if (sessionId === currentSessionId) return; + setOptimisticSessionId(sessionId); + setUrlSessionId(sessionId, { shallow: false }); + if (isMobile) handleCloseDrawer(); + } + + async function performNewChat() { + const sourceSessionId = currentSessionId; + + if (sourceSessionId) { + setIsSwitchingSession(true); + + await new Promise(function waitForStreamComplete(resolve) { + const unsubscribe = onStreamComplete( + function handleComplete(completedId) { + if (completedId === sourceSessionId) { + clearTimeout(timeout); + unsubscribe(); + resolve(); + } + }, + ); + const timeout = setTimeout(function handleTimeout() { + unsubscribe(); + resolve(); + }, 3000); + stopStream(sourceSessionId); + }); + + queryClient.invalidateQueries({ + queryKey: getGetV2GetSessionQueryKey(sourceSessionId), + }); + setIsSwitchingSession(false); + } + + resetAutoSelect(); + resetPagination(); + queryClient.invalidateQueries({ + queryKey: getGetV2ListSessionsQueryKey(), + }); + setUrlSessionId(null, { shallow: false }); + setOptimisticSessionId(null); if (isMobile) handleCloseDrawer(); } @@ -167,6 +266,7 @@ export function useCopilotShell() { queryKey: getGetV2ListSessionsQueryKey(), }); setUrlSessionId(null, { shallow: false }); + setOptimisticSessionId(null); if (isMobile) handleCloseDrawer(); } @@ -187,10 +287,12 @@ export function useCopilotShell() { sessions: visibleSessions, currentSessionId: sidebarSelectedSessionId, handleSelectSession, + performSelectSession, handleOpenDrawer, handleCloseDrawer, handleDrawerOpenChange, handleNewChat, + performNewChat, hasNextPage, isFetchingNextPage: isSessionsFetching, fetchNextPage, diff --git a/autogpt_platform/frontend/src/app/(platform)/copilot/copilot-page-store.ts b/autogpt_platform/frontend/src/app/(platform)/copilot/copilot-page-store.ts index 22bf5000a1..486d31865b 100644 --- a/autogpt_platform/frontend/src/app/(platform)/copilot/copilot-page-store.ts +++ b/autogpt_platform/frontend/src/app/(platform)/copilot/copilot-page-store.ts @@ -4,51 +4,106 @@ import { create } from "zustand"; interface CopilotStoreState { isStreaming: boolean; - isNewChatModalOpen: boolean; + isSwitchingSession: boolean; + isInterruptModalOpen: boolean; + pendingAction: (() => void) | null; newChatHandler: (() => void) | null; + newChatWithInterruptHandler: (() => void) | null; + selectSessionHandler: ((sessionId: string) => void) | null; + selectSessionWithInterruptHandler: ((sessionId: string) => void) | null; } interface CopilotStoreActions { setIsStreaming: (isStreaming: boolean) => void; + setIsSwitchingSession: (isSwitchingSession: boolean) => void; setNewChatHandler: (handler: (() => void) | null) => void; + setNewChatWithInterruptHandler: (handler: (() => void) | null) => void; + setSelectSessionHandler: ( + handler: ((sessionId: string) => void) | null, + ) => void; + setSelectSessionWithInterruptHandler: ( + handler: ((sessionId: string) => void) | null, + ) => void; requestNewChat: () => void; - confirmNewChat: () => void; - cancelNewChat: () => void; + requestSelectSession: (sessionId: string) => void; + confirmInterrupt: () => void; + cancelInterrupt: () => void; } type CopilotStore = CopilotStoreState & CopilotStoreActions; export const useCopilotStore = create((set, get) => ({ isStreaming: false, - isNewChatModalOpen: false, + isSwitchingSession: false, + isInterruptModalOpen: false, + pendingAction: null, newChatHandler: null, + newChatWithInterruptHandler: null, + selectSessionHandler: null, + selectSessionWithInterruptHandler: null, setIsStreaming(isStreaming) { set({ isStreaming }); }, + setIsSwitchingSession(isSwitchingSession) { + set({ isSwitchingSession }); + }, + setNewChatHandler(handler) { set({ newChatHandler: handler }); }, + setNewChatWithInterruptHandler(handler) { + set({ newChatWithInterruptHandler: handler }); + }, + + setSelectSessionHandler(handler) { + set({ selectSessionHandler: handler }); + }, + + setSelectSessionWithInterruptHandler(handler) { + set({ selectSessionWithInterruptHandler: handler }); + }, + requestNewChat() { - const { isStreaming, newChatHandler } = get(); + const { isStreaming, newChatHandler, newChatWithInterruptHandler } = get(); if (isStreaming) { - set({ isNewChatModalOpen: true }); + if (!newChatWithInterruptHandler) return; + set({ + isInterruptModalOpen: true, + pendingAction: newChatWithInterruptHandler, + }); } else if (newChatHandler) { newChatHandler(); } }, - confirmNewChat() { - const { newChatHandler } = get(); - set({ isNewChatModalOpen: false }); - if (newChatHandler) { - newChatHandler(); + requestSelectSession(sessionId) { + const { + isStreaming, + selectSessionHandler, + selectSessionWithInterruptHandler, + } = get(); + if (isStreaming) { + if (!selectSessionWithInterruptHandler) return; + set({ + isInterruptModalOpen: true, + pendingAction: () => selectSessionWithInterruptHandler(sessionId), + }); + } else { + if (!selectSessionHandler) return; + selectSessionHandler(sessionId); } }, - cancelNewChat() { - set({ isNewChatModalOpen: false }); + confirmInterrupt() { + const { pendingAction } = get(); + set({ isInterruptModalOpen: false, pendingAction: null }); + if (pendingAction) pendingAction(); + }, + + cancelInterrupt() { + set({ isInterruptModalOpen: false, pendingAction: null }); }, })); diff --git a/autogpt_platform/frontend/src/app/(platform)/copilot/page.tsx b/autogpt_platform/frontend/src/app/(platform)/copilot/page.tsx index 83b21bf82e..008b06fcda 100644 --- a/autogpt_platform/frontend/src/app/(platform)/copilot/page.tsx +++ b/autogpt_platform/frontend/src/app/(platform)/copilot/page.tsx @@ -13,22 +13,15 @@ import { useCopilotPage } from "./useCopilotPage"; export default function CopilotPage() { const { state, handlers } = useCopilotPage(); - const confirmNewChat = useCopilotStore((s) => s.confirmNewChat); - const { - greetingName, - quickActions, - isLoading, - pageState, - isNewChatModalOpen, - isReady, - } = state; + const isInterruptModalOpen = useCopilotStore((s) => s.isInterruptModalOpen); + const confirmInterrupt = useCopilotStore((s) => s.confirmInterrupt); + const cancelInterrupt = useCopilotStore((s) => s.cancelInterrupt); + const { greetingName, quickActions, isLoading, pageState, isReady } = state; const { handleQuickAction, startChatWithPrompt, handleSessionNotFound, handleStreamingChange, - handleCancelNewChat, - handleNewChatModalOpen, } = handlers; if (!isReady) return null; @@ -48,31 +41,33 @@ export default function CopilotPage() { title="Interrupt current chat?" styling={{ maxWidth: 300, width: "100%" }} controlled={{ - isOpen: isNewChatModalOpen, - set: handleNewChatModalOpen, + isOpen: isInterruptModalOpen, + set: (open) => { + if (!open) cancelInterrupt(); + }, }} - onClose={handleCancelNewChat} + onClose={cancelInterrupt} >
The current chat response will be interrupted. Are you sure you - want to start a new chat? + want to continue?
diff --git a/autogpt_platform/frontend/src/app/(platform)/copilot/useCopilotPage.ts b/autogpt_platform/frontend/src/app/(platform)/copilot/useCopilotPage.ts index 1d9c843d7d..8cf4599a12 100644 --- a/autogpt_platform/frontend/src/app/(platform)/copilot/useCopilotPage.ts +++ b/autogpt_platform/frontend/src/app/(platform)/copilot/useCopilotPage.ts @@ -75,9 +75,7 @@ export function useCopilotPage() { const { user, isLoggedIn, isUserLoading } = useSupabase(); const { toast } = useToast(); - const isNewChatModalOpen = useCopilotStore((s) => s.isNewChatModalOpen); const setIsStreaming = useCopilotStore((s) => s.setIsStreaming); - const cancelNewChat = useCopilotStore((s) => s.cancelNewChat); const isChatEnabled = useGetFlag(Flag.CHAT); const flags = useFlags(); @@ -201,21 +199,12 @@ export function useCopilotPage() { setIsStreaming(isStreamingValue); } - function handleCancelNewChat() { - cancelNewChat(); - } - - function handleNewChatModalOpen(isOpen: boolean) { - if (!isOpen) cancelNewChat(); - } - return { state: { greetingName, quickActions, isLoading: isUserLoading, pageState: state.pageState, - isNewChatModalOpen, isReady: isFlagReady && isChatEnabled !== false && isLoggedIn, }, handlers: { @@ -223,8 +212,6 @@ export function useCopilotPage() { startChatWithPrompt, handleSessionNotFound, handleStreamingChange, - handleCancelNewChat, - handleNewChatModalOpen, }, }; } diff --git a/autogpt_platform/frontend/src/components/contextual/Chat/Chat.tsx b/autogpt_platform/frontend/src/components/contextual/Chat/Chat.tsx index ba7584765d..a7a5f61674 100644 --- a/autogpt_platform/frontend/src/components/contextual/Chat/Chat.tsx +++ b/autogpt_platform/frontend/src/components/contextual/Chat/Chat.tsx @@ -1,11 +1,12 @@ "use client"; +import { useCopilotStore } from "@/app/(platform)/copilot/copilot-page-store"; +import { LoadingSpinner } from "@/components/atoms/LoadingSpinner/LoadingSpinner"; import { Text } from "@/components/atoms/Text/Text"; import { cn } from "@/lib/utils"; import { useEffect, useRef } from "react"; import { ChatContainer } from "./components/ChatContainer/ChatContainer"; import { ChatErrorState } from "./components/ChatErrorState/ChatErrorState"; -import { ChatLoader } from "./components/ChatLoader/ChatLoader"; import { useChat } from "./useChat"; export interface ChatProps { @@ -24,6 +25,7 @@ export function Chat({ onStreamingChange, }: ChatProps) { const hasHandledNotFoundRef = useRef(false); + const isSwitchingSession = useCopilotStore((s) => s.isSwitchingSession); const { messages, isLoading, @@ -47,29 +49,34 @@ export function Chat({ [onSessionNotFound, urlSessionId, isSessionNotFound, isLoading, isCreating], ); + const shouldShowLoader = + (showLoader && (isLoading || isCreating)) || isSwitchingSession; + return (
{/* Main Content */}
{/* Loading State */} - {showLoader && (isLoading || isCreating) && ( + {shouldShowLoader && (
-
- +
+ - Loading your chats... + {isSwitchingSession + ? "Switching chat..." + : "Loading your chat..."}
)} {/* Error State */} - {error && !isLoading && ( + {error && !isLoading && !isSwitchingSession && ( )} {/* Session Content */} - {sessionId && !isLoading && !error && ( + {sessionId && !isLoading && !error && !isSwitchingSession && ( ) { +function cleanupExpiredStreams( + completedStreams: Map, +): Map { const now = Date.now(); - for (const [sessionId, result] of completedStreams) { + const cleaned = new Map(completedStreams); + for (const [sessionId, result] of cleaned) { if (now - result.completedAt > COMPLETED_STREAM_TTL) { - completedStreams.delete(sessionId); + cleaned.delete(sessionId); } } -} - -function moveToCompleted( - activeStreams: Map, - completedStreams: Map, - streamCompleteCallbacks: Set, - sessionId: string, -) { - const stream = activeStreams.get(sessionId); - if (!stream) return; - - const result: StreamResult = { - sessionId, - status: stream.status, - chunks: stream.chunks, - completedAt: Date.now(), - error: stream.error, - }; - - completedStreams.set(sessionId, result); - activeStreams.delete(sessionId); - cleanupCompletedStreams(completedStreams); - - if (stream.status === "completed" || stream.status === "error") { - notifyStreamComplete(streamCompleteCallbacks, sessionId); - } + return cleaned; } export const useChatStore = create((set, get) => ({ @@ -106,17 +84,31 @@ export const useChatStore = create((set, get) => ({ context, onChunk, ) { - const { activeStreams, completedStreams, streamCompleteCallbacks } = get(); + const state = get(); + const newActiveStreams = new Map(state.activeStreams); + let newCompletedStreams = new Map(state.completedStreams); + const callbacks = state.streamCompleteCallbacks; - const existingStream = activeStreams.get(sessionId); + const existingStream = newActiveStreams.get(sessionId); if (existingStream) { existingStream.abortController.abort(); - moveToCompleted( - activeStreams, - completedStreams, - streamCompleteCallbacks, + const normalizedStatus = + existingStream.status === "streaming" + ? "completed" + : existingStream.status; + const result: StreamResult = { sessionId, - ); + status: normalizedStatus, + chunks: existingStream.chunks, + completedAt: Date.now(), + error: existingStream.error, + }; + newCompletedStreams.set(sessionId, result); + newActiveStreams.delete(sessionId); + newCompletedStreams = cleanupExpiredStreams(newCompletedStreams); + if (normalizedStatus === "completed" || normalizedStatus === "error") { + notifyStreamComplete(callbacks, sessionId); + } } const abortController = new AbortController(); @@ -132,36 +124,76 @@ export const useChatStore = create((set, get) => ({ onChunkCallbacks: initialCallbacks, }; - activeStreams.set(sessionId, stream); + newActiveStreams.set(sessionId, stream); + set({ + activeStreams: newActiveStreams, + completedStreams: newCompletedStreams, + }); try { await executeStream(stream, message, isUserMessage, context); } finally { if (onChunk) stream.onChunkCallbacks.delete(onChunk); if (stream.status !== "streaming") { - moveToCompleted( - activeStreams, - completedStreams, - streamCompleteCallbacks, - sessionId, - ); + const currentState = get(); + const finalActiveStreams = new Map(currentState.activeStreams); + let finalCompletedStreams = new Map(currentState.completedStreams); + + const storedStream = finalActiveStreams.get(sessionId); + if (storedStream === stream) { + const result: StreamResult = { + sessionId, + status: stream.status, + chunks: stream.chunks, + completedAt: Date.now(), + error: stream.error, + }; + finalCompletedStreams.set(sessionId, result); + finalActiveStreams.delete(sessionId); + finalCompletedStreams = cleanupExpiredStreams(finalCompletedStreams); + set({ + activeStreams: finalActiveStreams, + completedStreams: finalCompletedStreams, + }); + if (stream.status === "completed" || stream.status === "error") { + notifyStreamComplete( + currentState.streamCompleteCallbacks, + sessionId, + ); + } + } } } }, stopStream: function stopStream(sessionId) { - const { activeStreams, completedStreams, streamCompleteCallbacks } = get(); - const stream = activeStreams.get(sessionId); - if (stream) { - stream.abortController.abort(); - stream.status = "completed"; - moveToCompleted( - activeStreams, - completedStreams, - streamCompleteCallbacks, - sessionId, - ); - } + const state = get(); + const stream = state.activeStreams.get(sessionId); + if (!stream) return; + + stream.abortController.abort(); + stream.status = "completed"; + + const newActiveStreams = new Map(state.activeStreams); + let newCompletedStreams = new Map(state.completedStreams); + + const result: StreamResult = { + sessionId, + status: stream.status, + chunks: stream.chunks, + completedAt: Date.now(), + error: stream.error, + }; + newCompletedStreams.set(sessionId, result); + newActiveStreams.delete(sessionId); + newCompletedStreams = cleanupExpiredStreams(newCompletedStreams); + + set({ + activeStreams: newActiveStreams, + completedStreams: newCompletedStreams, + }); + + notifyStreamComplete(state.streamCompleteCallbacks, sessionId); }, subscribeToStream: function subscribeToStream( @@ -169,16 +201,18 @@ export const useChatStore = create((set, get) => ({ onChunk, skipReplay = false, ) { - const { activeStreams } = get(); + const state = get(); + const stream = state.activeStreams.get(sessionId); - const stream = activeStreams.get(sessionId); if (stream) { if (!skipReplay) { for (const chunk of stream.chunks) { onChunk(chunk); } } + stream.onChunkCallbacks.add(onChunk); + return function unsubscribe() { stream.onChunkCallbacks.delete(onChunk); }; @@ -204,7 +238,12 @@ export const useChatStore = create((set, get) => ({ }, clearCompletedStream: function clearCompletedStream(sessionId) { - get().completedStreams.delete(sessionId); + const state = get(); + if (!state.completedStreams.has(sessionId)) return; + + const newCompletedStreams = new Map(state.completedStreams); + newCompletedStreams.delete(sessionId); + set({ completedStreams: newCompletedStreams }); }, isStreaming: function isStreaming(sessionId) { @@ -213,11 +252,21 @@ export const useChatStore = create((set, get) => ({ }, registerActiveSession: function registerActiveSession(sessionId) { - get().activeSessions.add(sessionId); + const state = get(); + if (state.activeSessions.has(sessionId)) return; + + const newActiveSessions = new Set(state.activeSessions); + newActiveSessions.add(sessionId); + set({ activeSessions: newActiveSessions }); }, unregisterActiveSession: function unregisterActiveSession(sessionId) { - get().activeSessions.delete(sessionId); + const state = get(); + if (!state.activeSessions.has(sessionId)) return; + + const newActiveSessions = new Set(state.activeSessions); + newActiveSessions.delete(sessionId); + set({ activeSessions: newActiveSessions }); }, isSessionActive: function isSessionActive(sessionId) { @@ -225,10 +274,16 @@ export const useChatStore = create((set, get) => ({ }, onStreamComplete: function onStreamComplete(callback) { - const { streamCompleteCallbacks } = get(); - streamCompleteCallbacks.add(callback); + const state = get(); + const newCallbacks = new Set(state.streamCompleteCallbacks); + newCallbacks.add(callback); + set({ streamCompleteCallbacks: newCallbacks }); + return function unsubscribe() { - streamCompleteCallbacks.delete(callback); + const currentState = get(); + const cleanedCallbacks = new Set(currentState.streamCompleteCallbacks); + cleanedCallbacks.delete(callback); + set({ streamCompleteCallbacks: cleanedCallbacks }); }; }, })); diff --git a/autogpt_platform/frontend/src/components/contextual/Chat/components/ChatMessage/ChatMessage.tsx b/autogpt_platform/frontend/src/components/contextual/Chat/components/ChatMessage/ChatMessage.tsx index 0fee33dbc0..29e3a60a8c 100644 --- a/autogpt_platform/frontend/src/components/contextual/Chat/components/ChatMessage/ChatMessage.tsx +++ b/autogpt_platform/frontend/src/components/contextual/Chat/components/ChatMessage/ChatMessage.tsx @@ -126,10 +126,6 @@ export function ChatMessage({ [displayContent, message], ); - function isLongResponse(content: string): boolean { - return content.split("\n").length > 5; - } - const handleTryAgain = useCallback(() => { if (message.type !== "message" || !onSendMessage) return; onSendMessage(message.content, message.role === "user"); @@ -358,7 +354,7 @@ export function ChatMessage({ )} - {!isUser && isFinalMessage && isLongResponse(displayContent) && ( + {!isUser && isFinalMessage && !isStreaming && (