diff --git a/ui/web/src/pages/chat/chat-page.tsx b/ui/web/src/pages/chat/chat-page.tsx index 36fbf612..5df229f8 100644 --- a/ui/web/src/pages/chat/chat-page.tsx +++ b/ui/web/src/pages/chat/chat-page.tsx @@ -81,8 +81,8 @@ export function ChatPage() { const isOwn = !sessionKey || isOwnSession(sessionKey, userId); const handleMessageAdded = useCallback( - (msg: { role: "user" | "assistant" | "tool"; content: string; timestamp?: number }) => { - addLocalMessage(msg); + (msg: { role: "user" | "assistant" | "tool"; content: string; timestamp?: number }, key?: string) => { + addLocalMessage(msg, key); }, [addLocalMessage], ); diff --git a/ui/web/src/pages/chat/hooks/use-chat-messages.ts b/ui/web/src/pages/chat/hooks/use-chat-messages.ts index 5f11f15a..f59fa420 100644 --- a/ui/web/src/pages/chat/hooks/use-chat-messages.ts +++ b/ui/web/src/pages/chat/hooks/use-chat-messages.ts @@ -7,6 +7,11 @@ import type { ChatMessage, AgentEventPayload, ToolStreamEntry, RunActivity, Acti import { toFileUrl, mediaKindFromMime } from "@/lib/file-helpers"; import { transformHistoryMessages } from "@/adapters/chat-message.adapter"; import { useChatTeamTasks } from "./use-chat-team-tasks"; +import { useChatMessagesStore } from "@/stores/use-chat-messages-store"; + +// Stable empty array — avoids creating a new reference on every render inside +// Zustand selectors, which would trigger an infinite re-render loop (React #185). +const EMPTY_MESSAGES: ChatMessage[] = []; /** * Manages chat message history and real-time streaming for a session. @@ -14,11 +19,19 @@ import { useChatTeamTasks } from "./use-chat-team-tasks"; */ export function useChatMessages(sessionKey: string, agentId: string) { const ws = useWs(); - const [messages, setMessages] = useState([]); - const [streamText, setStreamText] = useState(null); - const [thinkingText, setThinkingText] = useState(null); + const messages = useChatMessagesStore((s) => sessionKey ? (s.sessions[sessionKey]?.messages ?? EMPTY_MESSAGES) : EMPTY_MESSAGES); + const streamText = useChatMessagesStore((s) => sessionKey ? (s.sessions[sessionKey]?.streamText ?? null) : null); + const thinkingText = useChatMessagesStore((s) => sessionKey ? (s.sessions[sessionKey]?.thinkingText ?? null) : null); + const isRunning = useChatMessagesStore((s) => sessionKey ? (s.sessions[sessionKey]?.isRunning ?? false) : false); + + const setSessionMessages = useChatMessagesStore((s) => s.setSessionMessages); + const updateSessionMessages = useChatMessagesStore((s) => s.updateSessionMessages); + const setSessionStream = useChatMessagesStore((s) => s.setSessionStream); + const setSessionThinking = useChatMessagesStore((s) => s.setSessionThinking); + const setSessionRunning = useChatMessagesStore((s) => s.setSessionRunning); + + // Local state for non-persistent UI elements const [toolStream, setToolStream] = useState([]); - const [isRunning, setIsRunning] = useState(false); const [loading, setLoading] = useState(false); const [activity, setActivity] = useState(null); const [blockReplies, setBlockReplies] = useState([]); @@ -38,10 +51,15 @@ export function useChatMessages(sessionKey: string, agentId: string) { const rafPendingRef = useRef(false); const rafHandleRef = useRef(0); - // Add a local message optimistically - const addLocalMessage = useCallback((msg: ChatMessage) => { - setMessages((prev) => [...prev, msg]); - }, []); + // Add a local message optimistically. + // `key` is optional: callers that know the target session key (e.g. new-chat + // send flow, where the URL hasn't navigated yet) should pass it explicitly. + const addLocalMessage = useCallback((msg: ChatMessage, key?: string) => { + const targetKey = key ?? sessionKey; + if (targetKey) { + updateSessionMessages(targetKey, (prev) => [...prev, msg]); + } + }, [sessionKey, updateSessionMessages]); // Team task handling (extracted hook) const { teamTasks, setTeamTasks } = useChatTeamTasks(addLocalMessage); @@ -67,23 +85,32 @@ export function useChatMessages(sessionKey: string, agentId: string) { return; } - setStreamText(null); setThinkingText(null); setToolStream([]); - setIsRunning(false); setActivity(null); setBlockReplies([]); setTeamTasks([]); - runIdRef.current = null; expectingRunRef.current = false; - streamRef.current = ""; thinkingRef.current = ""; - toolStreamRef.current = []; activityRef.current = null; blockRepliesRef.current = []; - cancelAnimationFrame(rafHandleRef.current); rafPendingRef.current = false; - if (!sessionKey) setMessages([]); - }, [sessionKey, setTeamTasks]); + setSessionStream(sessionKey, null); + setSessionThinking(sessionKey, null); + setSessionRunning(sessionKey, false); + setToolStream([]); + setActivity(null); + setBlockReplies([]); + setTeamTasks([]); + runIdRef.current = null; + expectingRunRef.current = false; + streamRef.current = ""; + thinkingRef.current = ""; + toolStreamRef.current = []; + activityRef.current = null; + blockRepliesRef.current = []; + cancelAnimationFrame(rafHandleRef.current); + rafPendingRef.current = false; + }, [sessionKey, setTeamTasks, setSessionStream, setSessionThinking, setSessionRunning]); // Load history const loadHistory = useCallback(async (mediaItems?: MediaItem[]) => { if (!ws.isConnected || !sessionKey) { setLoading(false); return; } try { const res = await ws.call<{ messages: Message[] }>(Methods.CHAT_HISTORY, { agentId, sessionKey }); - setMessages(transformHistoryMessages(res.messages ?? [], mediaItems)); + setSessionMessages(sessionKey, transformHistoryMessages(res.messages ?? [], mediaItems)); } catch { /* will retry */ } finally { setLoading(false); } - }, [ws, agentId, sessionKey]); + }, [ws, agentId, sessionKey, setSessionMessages]); // Load history + restore running state when session changes useEffect(() => { @@ -99,7 +126,7 @@ export function useChatMessages(sessionKey: string, agentId: string) { ws.call<{ isRunning?: boolean; runId?: string; activity?: RunActivity }>(Methods.CHAT_SESSION_STATUS, { sessionKey }) .then((res) => { if (cancelled) return; - if (res.isRunning) { setIsRunning(true); if (res.runId) runIdRef.current = res.runId; } + if (res.isRunning) { setSessionRunning(sessionKey, true); if (res.runId) runIdRef.current = res.runId; } if (res.activity) { setActivity(res.activity); activityRef.current = res.activity; } }).catch((err) => console.error("[useChatMessages] session status failed:", err)); ws.call<{ tasks?: ActiveTeamTask[] }>(Methods.TEAMS_TASK_ACTIVE_BY_SESSION, { sessionKey }) @@ -107,7 +134,7 @@ export function useChatMessages(sessionKey: string, agentId: string) { .catch((err) => console.error("[useChatMessages] active tasks failed:", err)); } return () => { cancelled = true; }; - }, [sessionKey, loadHistory, ws, setTeamTasks]); + }, [sessionKey, loadHistory, ws, setTeamTasks, setSessionRunning]); // Called before sending so event handler captures run.started const expectRun = useCallback(() => { expectingRunRef.current = true; }, []); @@ -123,9 +150,15 @@ export function useChatMessages(sessionKey: string, agentId: string) { // Capture run.started if (event.type === "run.started" && event.agentId === agentIdRef.current) { if (expectingRunRef.current || event.runKind === "announce") { - runIdRef.current = event.runId; expectingRunRef.current = false; - setIsRunning(true); setStreamText(null); setThinkingText(null); setToolStream([]); - streamRef.current = ""; thinkingRef.current = ""; toolStreamRef.current = []; + runIdRef.current = event.runId; + expectingRunRef.current = false; + setSessionRunning(sessionKeyRef.current, true); + setSessionStream(sessionKeyRef.current, null); + setSessionThinking(sessionKeyRef.current, null); + setToolStream([]); + streamRef.current = ""; + thinkingRef.current = ""; + toolStreamRef.current = []; } return; } @@ -139,7 +172,8 @@ export function useChatMessages(sessionKey: string, agentId: string) { rafPendingRef.current = true; rafHandleRef.current = requestAnimationFrame(() => { rafPendingRef.current = false; - setThinkingText(thinkingRef.current); setStreamText(streamRef.current); + setSessionThinking(sessionKeyRef.current, thinkingRef.current); + setSessionStream(sessionKeyRef.current, streamRef.current); }); } break; @@ -150,7 +184,8 @@ export function useChatMessages(sessionKey: string, agentId: string) { rafPendingRef.current = true; rafHandleRef.current = requestAnimationFrame(() => { rafPendingRef.current = false; - setStreamText(streamRef.current); setThinkingText(thinkingRef.current); + setSessionStream(sessionKeyRef.current, streamRef.current); + setSessionThinking(sessionKeyRef.current, thinkingRef.current); }); } break; @@ -201,48 +236,69 @@ export function useChatMessages(sessionKey: string, agentId: string) { } case "run.completed": { cancelAnimationFrame(rafHandleRef.current); rafPendingRef.current = false; - setIsRunning(false); runIdRef.current = null; + setSessionRunning(sessionKeyRef.current, false); + runIdRef.current = null; const hadTools = toolStreamRef.current.length > 0; const streamed = streamRef.current; const thinking = thinkingRef.current || undefined; - setStreamText(null); setThinkingText(null); setToolStream([]); - streamRef.current = ""; thinkingRef.current = ""; toolStreamRef.current = []; - activityRef.current = null; setActivity(null); - blockRepliesRef.current = []; setBlockReplies([]); + setSessionStream(sessionKeyRef.current, null); + setSessionThinking(sessionKeyRef.current, null); + setToolStream([]); + streamRef.current = ""; + thinkingRef.current = ""; + toolStreamRef.current = []; + activityRef.current = null; + setActivity(null); + blockRepliesRef.current = []; + setBlockReplies([]); const rawMedia = event.payload?.media; const mediaItems: MediaItem[] | undefined = rawMedia?.length ? rawMedia.map((m) => ({ path: toFileUrl(m.path), mimeType: m.content_type ?? "application/octet-stream", fileName: m.path.split("?")[0]?.split("/").pop() ?? "file", size: m.size, kind: mediaKindFromMime(m.content_type ?? "") })) : undefined; if (streamed && !hadTools) { - setMessages((prev) => [...prev, { role: "assistant", content: streamed, thinking, timestamp: Date.now(), mediaItems }]); + updateSessionMessages(sessionKeyRef.current, (prev) => [...prev, { role: "assistant", content: streamed, thinking, timestamp: Date.now(), mediaItems }]); } else { loadHistory(mediaItems); } break; } case "run.failed": { cancelAnimationFrame(rafHandleRef.current); rafPendingRef.current = false; - setIsRunning(false); runIdRef.current = null; - setStreamText(null); setThinkingText(null); setToolStream([]); - streamRef.current = ""; thinkingRef.current = ""; - activityRef.current = null; setActivity(null); - blockRepliesRef.current = []; setBlockReplies([]); - setMessages((prev) => [...prev, { role: "assistant", content: `Error: ${event.payload?.error ?? "Unknown error"}`, timestamp: Date.now() }]); + setSessionRunning(sessionKeyRef.current, false); + runIdRef.current = null; + setSessionStream(sessionKeyRef.current, null); + setSessionThinking(sessionKeyRef.current, null); + setToolStream([]); + streamRef.current = ""; + thinkingRef.current = ""; + activityRef.current = null; + setActivity(null); + blockRepliesRef.current = []; + setBlockReplies([]); + updateSessionMessages(sessionKeyRef.current, (prev) => [...prev, { role: "assistant", content: `Error: ${event.payload?.error ?? "Unknown error"}`, timestamp: Date.now() }]); break; } case "run.cancelled": { cancelAnimationFrame(rafHandleRef.current); rafPendingRef.current = false; - setIsRunning(false); runIdRef.current = null; + setSessionRunning(sessionKeyRef.current, false); + runIdRef.current = null; const streamed = streamRef.current; - setStreamText(null); setThinkingText(null); setToolStream([]); - streamRef.current = ""; thinkingRef.current = ""; toolStreamRef.current = []; - activityRef.current = null; setActivity(null); - blockRepliesRef.current = []; setBlockReplies([]); - if (streamed) { setMessages((prev) => [...prev, { role: "assistant", content: streamed, timestamp: Date.now() }]); } - else { loadHistory(); } + setSessionStream(sessionKeyRef.current, null); + setSessionThinking(sessionKeyRef.current, null); + setToolStream([]); + streamRef.current = ""; + thinkingRef.current = ""; + toolStreamRef.current = []; + activityRef.current = null; + setActivity(null); + blockRepliesRef.current = []; + setBlockReplies([]); + if (streamed) { + updateSessionMessages(sessionKeyRef.current, (prev) => [...prev, { role: "assistant", content: streamed, timestamp: Date.now() }]); + } else { loadHistory(); } break; } } }, - [loadHistory], + [loadHistory, setSessionRunning, setSessionStream, setSessionThinking, updateSessionMessages], ); useWsEvent(Events.AGENT, handleAgentEvent); diff --git a/ui/web/src/pages/chat/hooks/use-chat-send.ts b/ui/web/src/pages/chat/hooks/use-chat-send.ts index 5bdc2d43..2e54d6ef 100644 --- a/ui/web/src/pages/chat/hooks/use-chat-send.ts +++ b/ui/web/src/pages/chat/hooks/use-chat-send.ts @@ -7,7 +7,7 @@ import type { AttachedFile } from "@/components/chat/chat-input"; interface UseChatSendOptions { agentId: string; - onMessageAdded: (msg: ChatMessage) => void; + onMessageAdded: (msg: ChatMessage, sessionKey?: string) => void; onExpectRun: () => void; } @@ -58,7 +58,7 @@ export function useChatSend({ role: "user", content: displayContent, timestamp: Date.now(), - }); + }, sessionKey); try { // Upload files first, then pass path+filename to chat.send diff --git a/ui/web/src/stores/use-chat-messages-store.ts b/ui/web/src/stores/use-chat-messages-store.ts new file mode 100644 index 00000000..98ee7239 --- /dev/null +++ b/ui/web/src/stores/use-chat-messages-store.ts @@ -0,0 +1,119 @@ +import { create } from "zustand"; +import type { ChatMessage } from "@/types/chat"; + +interface SessionMessages { + messages: ChatMessage[]; + streamText: string | null; + thinkingText: string | null; + isRunning: boolean; +} + +interface ChatMessagesState { + // Map of sessionKey -> session message state + sessions: Record; + + // Set messages for a specific session + setSessionMessages: (sessionKey: string, messages: ChatMessage[]) => void; + + // Update session messages incrementally + updateSessionMessages: (sessionKey: string, updater: (prev: ChatMessage[]) => ChatMessage[]) => void; + + // Set streaming text for a session + setSessionStream: (sessionKey: string, streamText: string | null) => void; + + // Set thinking text for a session + setSessionThinking: (sessionKey: string, thinkingText: string | null) => void; + + // Set running state for a session + setSessionRunning: (sessionKey: string, isRunning: boolean) => void; +} + +export const useChatMessagesStore = create((set) => ({ + sessions: {}, + + setSessionMessages: (sessionKey: string, messages: ChatMessage[]) => { + set((state) => { + const existing = state.sessions[sessionKey]; + return { + sessions: { + ...state.sessions, + [sessionKey]: { + messages, + streamText: existing?.streamText ?? null, + thinkingText: existing?.thinkingText ?? null, + isRunning: existing?.isRunning ?? false, + }, + }, + }; + }); + }, + + updateSessionMessages: (sessionKey: string, updater: (prev: ChatMessage[]) => ChatMessage[]) => { + set((state) => { + const existing = state.sessions[sessionKey]; + const current = existing?.messages ?? []; + return { + sessions: { + ...state.sessions, + [sessionKey]: { + messages: updater(current), + streamText: existing?.streamText ?? null, + thinkingText: existing?.thinkingText ?? null, + isRunning: existing?.isRunning ?? false, + }, + }, + }; + }); + }, + + setSessionStream: (sessionKey: string, streamText: string | null) => { + set((state) => { + const existing = state.sessions[sessionKey]; + return { + sessions: { + ...state.sessions, + [sessionKey]: { + messages: existing?.messages ?? [], + streamText, + thinkingText: existing?.thinkingText ?? null, + isRunning: existing?.isRunning ?? false, + }, + }, + }; + }); + }, + + setSessionThinking: (sessionKey: string, thinkingText: string | null) => { + set((state) => { + const existing = state.sessions[sessionKey]; + return { + sessions: { + ...state.sessions, + [sessionKey]: { + messages: existing?.messages ?? [], + streamText: existing?.streamText ?? null, + thinkingText, + isRunning: existing?.isRunning ?? false, + }, + }, + }; + }); + }, + + setSessionRunning: (sessionKey: string, isRunning: boolean) => { + set((state) => { + const existing = state.sessions[sessionKey]; + return { + sessions: { + ...state.sessions, + [sessionKey]: { + messages: existing?.messages ?? [], + streamText: existing?.streamText ?? null, + thinkingText: existing?.thinkingText ?? null, + isRunning, + }, + }, + }; + }); + }, +}));