diff --git a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx index da134c4cf..10abe13b1 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx +++ b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx @@ -77,7 +77,10 @@ import { type ThreadListResponse, type ThreadRecord, } from "@/lib/chat/thread-persistence"; -import { extractUserTurnForNewChatApi } from "@/lib/chat/user-turn-api-parts"; +import { + extractUserTurnForNewChatApi, + type NewChatUserImagePayload, +} from "@/lib/chat/user-turn-api-parts"; import { NotFoundError } from "@/lib/error"; import { trackChatCreated, @@ -1337,15 +1340,24 @@ export default function NewChatPage() { * Handle regeneration (edit or reload) by calling the regenerate endpoint * and streaming the response. This rewinds the LangGraph checkpointer state. * - * @param newUserQuery - The new user query (for edit). Pass null/undefined for reload. + * @param newUserQuery - `null` = reload with same turn from the server. A string = edit + * (including an empty string when the edited turn is images-only); pass `editExtras` for images/content. */ const handleRegenerate = useCallback( - async (newUserQuery?: string | null) => { + async ( + newUserQuery: string | null, + editExtras?: { + userMessageContent: ThreadMessageLike["content"]; + userImages: NewChatUserImagePayload[]; + } + ) => { if (!threadId) { toast.error("Cannot regenerate: no active chat thread"); return; } + const isEdit = newUserQuery !== null; + // Abort any previous streaming request if (abortControllerRef.current) { abortControllerRef.current.abort(); @@ -1359,11 +1371,11 @@ export default function NewChatPage() { } // Extract the original user query BEFORE removing messages (for reload mode) - let userQueryToDisplay = newUserQuery; + let userQueryToDisplay: string | undefined; let originalUserMessageContent: ThreadMessageLike["content"] | null = null; let originalUserMessageMetadata: ThreadMessageLike["metadata"] | undefined; - if (!newUserQuery) { + if (!isEdit) { // Reload mode - find and preserve the last user message content const lastUserMessage = [...messages].reverse().find((m) => m.role === "user"); if (lastUserMessage) { @@ -1377,6 +1389,8 @@ export default function NewChatPage() { } } } + } else { + userQueryToDisplay = newUserQuery; } // Remove the last two messages (user + assistant) from the UI immediately @@ -1412,11 +1426,13 @@ export default function NewChatPage() { const userMessage: ThreadMessageLike = { id: userMsgId, role: "user", - content: newUserQuery - ? [{ type: "text", text: newUserQuery }] + content: isEdit + ? (editExtras?.userMessageContent ?? [ + { type: "text", text: newUserQuery ?? "" }, + ]) : originalUserMessageContent || [{ type: "text", text: userQueryToDisplay || "" }], createdAt: new Date(), - metadata: newUserQuery ? undefined : originalUserMessageMetadata, + metadata: isEdit ? undefined : originalUserMessageMetadata, }; setMessages((prev) => [...prev, userMessage]); @@ -1433,20 +1449,24 @@ export default function NewChatPage() { try { const selection = await getAgentFilesystemSelection(); + const requestBody: Record = { + search_space_id: searchSpaceId, + user_query: newUserQuery, + disabled_tools: disabledTools.length > 0 ? disabledTools : undefined, + filesystem_mode: selection.filesystem_mode, + client_platform: selection.client_platform, + local_filesystem_mounts: selection.local_filesystem_mounts, + }; + if (isEdit) { + requestBody.user_images = editExtras?.userImages ?? []; + } const response = await fetch(getRegenerateUrl(threadId), { method: "POST", headers: { "Content-Type": "application/json", Authorization: `Bearer ${token}`, }, - body: JSON.stringify({ - search_space_id: searchSpaceId, - user_query: newUserQuery || null, - disabled_tools: disabledTools.length > 0 ? disabledTools : undefined, - filesystem_mode: selection.filesystem_mode, - client_platform: selection.client_platform, - local_filesystem_mounts: selection.local_filesystem_mounts, - }), + body: JSON.stringify(requestBody), signal: controller.signal, }); @@ -1536,8 +1556,10 @@ export default function NewChatPage() { if (contentParts.length > 0) { try { // Persist user message (for both edit and reload modes, since backend deleted it) - const userContentToPersist = newUserQuery - ? [{ type: "text", text: newUserQuery }] + const userContentToPersist = isEdit + ? (editExtras?.userMessageContent ?? [ + { type: "text", text: newUserQuery ?? "" }, + ]) : originalUserMessageContent || [{ type: "text", text: userQueryToDisplay || "" }]; const savedUserMessage = await appendMessage(threadId, { @@ -1602,21 +1624,15 @@ export default function NewChatPage() { // Handle editing a message - truncates history and regenerates with new query const onEdit = useCallback( async (message: AppendMessage) => { - // Extract the new user query from the message content - let newUserQuery = ""; - for (const part of message.content) { - if (part.type === "text") { - newUserQuery += part.text; - } - } - - if (!newUserQuery.trim()) { + const { userQuery, userImages } = extractUserTurnForNewChatApi(message, []); + const queryForApi = userQuery.trim(); + if (!queryForApi && userImages.length === 0) { toast.error("Cannot edit with empty message"); return; } - // Call regenerate with the new query - await handleRegenerate(newUserQuery.trim()); + const userMessageContent = message.content as unknown as ThreadMessageLike["content"]; + await handleRegenerate(queryForApi, { userMessageContent, userImages }); }, [handleRegenerate] );