From 204b3aa56cdcab49c6eaab3a27e3a019b8a8a90a Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Mon, 13 Jul 2026 22:17:38 +0200 Subject: [PATCH] feat(chat): add stream engine driving turns into the store --- .../lib/chat/stream-engine/engine.ts | 1433 +++++++++++++++++ 1 file changed, 1433 insertions(+) create mode 100644 surfsense_web/lib/chat/stream-engine/engine.ts diff --git a/surfsense_web/lib/chat/stream-engine/engine.ts b/surfsense_web/lib/chat/stream-engine/engine.ts new file mode 100644 index 000000000..a4a0c02a7 --- /dev/null +++ b/surfsense_web/lib/chat/stream-engine/engine.ts @@ -0,0 +1,1433 @@ +import type { AppendMessage, ThreadMessageLike } from "@assistant-ui/react"; +import { getDefaultStore } from "jotai"; +import type { Dispatch, SetStateAction } from "react"; +import { toast } from "sonner"; +import { agentFlagsAtom } from "@/atoms/agent/agent-flags-query.atom"; +import { disabledToolsAtom } from "@/atoms/agent-tools/agent-tools.atoms"; +import { + deriveMentionedPayload, + type MentionedDocumentInfo, + mentionedDocumentsAtom, + messageDocumentsMapAtom, + submittedMentionsAtom, +} from "@/atoms/chat/mentioned-documents.atom"; +import { pendingUserImageDataUrlsAtom } from "@/atoms/chat/pending-user-images.atom"; +import { setPremiumAlertForThreadAtom } from "@/atoms/chat/premium-alert.atom"; +import { type AgentCreatedDocument, agentCreatedDocumentsAtom } from "@/atoms/documents/ui.atoms"; +import { updateChatTabTitleAtom } from "@/atoms/tabs/tabs.atom"; +import { currentUserAtom } from "@/atoms/user/user-query.atoms"; +import type { HitlDecision, PendingInterruptState } from "@/features/chat-messages/hitl"; +import { + applyActionLogSse, + applyActionLogUpdatedSse, + markActionRevertedInCache, +} from "@/hooks/use-agent-actions-query"; +import { getAgentFilesystemSelection } from "@/lib/agent-filesystem"; +import { authenticatedFetch } from "@/lib/auth-fetch"; +import { type ChatFlow, classifyChatError } from "@/lib/chat/chat-error-classifier"; +import { tagPreAcceptSendFailure, toHttpResponseError } from "@/lib/chat/chat-request-errors"; +import { getMentionDocKey } from "@/lib/chat/mention-doc-key"; +import { createStreamFlushHelpers } from "@/lib/chat/stream-flush"; +import { consumeSseEvents, processSharedStreamEvent } from "@/lib/chat/stream-pipeline"; +import { + applyTurnIdToAssistantMessageList, + mergeChatTurnIdIntoMessage, + readStreamedChatTurnId, + readStreamedMessageId, +} from "@/lib/chat/stream-side-effects"; +import { + addToolCall, + buildContentForUI, + type ContentPartsState, + type FrameBatchedUpdater, + type ThinkingStepData, + updateToolCall, +} from "@/lib/chat/streaming-state"; +import { + appendMessage, + createThread, + getRegenerateUrl, + type ThreadListItem, + type ThreadListResponse, + type ThreadRecord, +} from "@/lib/chat/thread-persistence"; +import { + extractUserTurnForNewChatApi, + type NewChatUserImagePayload, +} from "@/lib/chat/user-turn-api-parts"; +import { buildBackendUrl } from "@/lib/env-config"; +import { + trackChatBlocked, + trackChatCreated, + trackChatErrorDetailed, + trackChatMessageSent, + trackChatResponseReceived, +} from "@/lib/posthog/events"; +import { cacheKeys } from "@/lib/query-client/cache-keys"; +import { queryClient } from "@/lib/query-client/client"; +import { + computeFallbackTurnCancellingRetryDelay, + freshSynthToolCallId, + pairBundleToolCallIds, + RECENT_CANCEL_WINDOW_MS, + sleep, + TOOLS_WITH_UI_ALL, +} from "./helpers"; +import { chatStreamStore } from "./store"; + +const jotaiStore = getDefaultStore(); +const tokenUsageStore = chatStreamStore.tokenUsage; +const toolsWithUI = TOOLS_WITH_UI_ALL; + +/** + * Display-only setters the page provides while it is mounted. After the chat + * page unmounts (in-app navigation) these are stale but harmless — the stream + * keeps writing to the module {@link chatStreamStore}, and the page re-derives + * its thread from the URL when it remounts. + */ +export interface EngineView { + setThreadId: Dispatch>; + setCurrentThread: Dispatch>; +} + +/** Route/view context the page passes into every engine call. */ +export interface EngineContext { + workspaceId: number; + /** Currently viewed thread id (``activeThreadId`` in the page). */ + threadId: number | null; + /** The page's current displayed messages — history/slice seed. */ + priorMessages: ThreadMessageLike[]; + view: EngineView; +} + +// --------------------------------------------------------------------------- +// Error handling (relocated from the chat page) +// --------------------------------------------------------------------------- + +async function persistAssistantErrorMessage({ + threadId, + assistantMsgId, + text, +}: { + threadId: number | null; + assistantMsgId: string; + text: string; +}): Promise { + if (threadId != null) { + chatStreamStore.setMessages(threadId, (prev) => + prev.map((m) => (m.id === assistantMsgId ? { ...m, content: [{ type: "text", text }] } : m)) + ); + } + + if (!threadId) return; + + // Persist only temporary assistant placeholders to avoid duplicate rows + // when the message already has a database-backed ID. + if (!assistantMsgId.startsWith("msg-assistant-")) return; + + try { + const savedMessage = await appendMessage(threadId, { + role: "assistant", + content: [{ type: "text", text }], + }); + const newMsgId = `msg-${savedMessage.id}`; + tokenUsageStore.rename(assistantMsgId, newMsgId); + chatStreamStore.setMessages(threadId, (prev) => + prev.map((m) => (m.id === assistantMsgId ? { ...m, id: newMsgId } : m)) + ); + } catch (persistErr) { + console.error("Failed to persist assistant error message:", persistErr); + } +} + +async function handleChatFailure({ + error, + flow, + threadId, + assistantMsgId, + workspaceId, +}: { + error: unknown; + flow: ChatFlow; + threadId: number | null; + assistantMsgId: string; + workspaceId: number; +}): Promise { + const normalized = classifyChatError({ + error, + flow, + context: { workspaceId, threadId }, + }); + + const logger = + normalized.severity === "error" + ? console.error + : normalized.severity === "warn" + ? console.warn + : console.info; + logger(`[chat-engine] ${flow} ${normalized.kind}:`, error); + + const telemetryPayload = { + flow, + kind: normalized.kind, + error_code: normalized.errorCode, + severity: normalized.severity, + is_expected: normalized.isExpected, + message: normalized.userMessage, + }; + if (normalized.telemetryEvent === "chat_blocked") { + trackChatBlocked(workspaceId, threadId, telemetryPayload); + } else { + trackChatErrorDetailed(workspaceId, threadId, telemetryPayload); + } + + if (normalized.channel === "silent") { + return; + } + + if (normalized.channel === "pinned_inline") { + if (threadId) { + const currentUser = jotaiStore.get(currentUserAtom).data; + jotaiStore.set(setPremiumAlertForThreadAtom, { + threadId, + message: normalized.userMessage, + userId: currentUser?.id ?? null, + }); + } + if (normalized.assistantMessage) { + await persistAssistantErrorMessage({ + threadId, + assistantMsgId, + text: normalized.assistantMessage, + }); + } + return; + } + + if (normalized.channel === "inline") { + if (normalized.assistantMessage) { + await persistAssistantErrorMessage({ + threadId, + assistantMsgId, + text: normalized.assistantMessage, + }); + } + toast.error(normalized.userMessage); + return; + } + + toast.error(normalized.userMessage); +} + +async function handleStreamTerminalError({ + error, + flow, + threadId, + assistantMsgId, + accepted, + workspaceId, + onAbort, + onPreAcceptFailure, + onAcceptedStreamError, +}: { + error: unknown; + flow: ChatFlow; + threadId: number | null; + assistantMsgId: string; + accepted: boolean; + workspaceId: number; + onAbort?: () => Promise; + onPreAcceptFailure?: () => Promise; + onAcceptedStreamError?: () => Promise; +}): Promise { + if (error instanceof Error && error.name === "AbortError") { + await onAbort?.(); + return; + } + + if (!accepted) { + await onPreAcceptFailure?.(); + } else { + await onAcceptedStreamError?.(); + } + + await handleChatFailure({ + error: !accepted ? tagPreAcceptSendFailure(error) : error, + flow, + threadId, + assistantMsgId: accepted ? assistantMsgId : "no-persist-assistant", + workspaceId, + }); +} + +async function fetchWithTurnCancellingRetry(runFetch: () => Promise): Promise { + const maxAttempts = 4; + for (let attempt = 1; attempt <= maxAttempts; attempt += 1) { + const response = await runFetch(); + if (response.ok) { + return response; + } + const error = await toHttpResponseError(response); + const withMeta = error as Error & { errorCode?: string; retryAfterMs?: number }; + const isTurnCancelling = withMeta.errorCode === "TURN_CANCELLING"; + const isRecentThreadBusyAfterCancel = + withMeta.errorCode === "THREAD_BUSY" && + Date.now() - chatStreamStore.recentCancelRequestedAt <= RECENT_CANCEL_WINDOW_MS; + if ((isTurnCancelling || isRecentThreadBusyAfterCancel) && attempt < maxAttempts) { + const waitMs = withMeta.retryAfterMs ?? computeFallbackTurnCancellingRetryDelay(attempt); + await sleep(waitMs); + continue; + } + throw error; + } + + throw Object.assign(new Error("Turn cancellation retry limit exceeded"), { + errorCode: "TURN_CANCELLING", + }); +} + +// --------------------------------------------------------------------------- +// Cancel +// --------------------------------------------------------------------------- + +/** + * Cancel the single in-flight turn. Targets the active stream's OWNER thread + * (not the currently-viewed thread) so cancel works even after navigation and + * for the lazy-create case. + */ +export async function cancelActiveTurn(): Promise { + const threadId = chatStreamStore.activeThreadId; + if (threadId) { + try { + const response = await authenticatedFetch( + buildBackendUrl(`/api/v1/threads/${threadId}/cancel-active-turn`), + { method: "POST" } + ); + if (response.ok) { + const payload = (await response.json()) as { error_code?: string }; + if (payload.error_code === "TURN_CANCELLING") { + chatStreamStore.markRecentCancel(); + } + } + } catch (error) { + console.warn("[chat-engine] Failed to signal cancel-active-turn:", error); + } + chatStreamStore.setRunning(threadId, false); + } + chatStreamStore.abortActive(); +} + +// --------------------------------------------------------------------------- +// New chat turn +// --------------------------------------------------------------------------- + +export async function startNewChat(ctx: EngineContext, message: AppendMessage): Promise { + const { workspaceId, threadId, priorMessages, view } = ctx; + + // Supersede any previous in-flight turn. + chatStreamStore.abortActive(); + + // Prefer the submit-time snapshot; fall back to the live atom for the + // send-button path. + const submittedSnapshot = jotaiStore.get(submittedMentionsAtom); + jotaiStore.set(submittedMentionsAtom, null); + const mentionedDocuments = jotaiStore.get(mentionedDocumentsAtom); + const activeMentions = submittedSnapshot ?? mentionedDocuments; + const mentionPayload = deriveMentionedPayload(activeMentions); + if (activeMentions.length > 0) { + jotaiStore.set(mentionedDocumentsAtom, []); + } + + const pendingUserImageUrls = jotaiStore.get(pendingUserImageDataUrlsAtom); + const urlsSnapshot = [...pendingUserImageUrls]; + const { userQuery, userImages } = extractUserTurnForNewChatApi(message, urlsSnapshot); + + if (!userQuery.trim() && userImages.length === 0) return; + + const localFilesystemEnabled = + jotaiStore.get(agentFlagsAtom).data?.enable_desktop_local_filesystem === true; + const disabledTools = jotaiStore.get(disabledToolsAtom); + const currentUser = jotaiStore.get(currentUserAtom).data; + + // Resolve filesystem selection BEFORE any optimistic UI / lazy thread + // creation so an unsatisfied "Local Folder" requirement bails cleanly + // with nothing to roll back and no thread left stuck "running". + let selection: Awaited>; + try { + selection = await getAgentFilesystemSelection(workspaceId, { localFilesystemEnabled }); + } catch (error) { + await handleChatFailure({ + error: tagPreAcceptSendFailure(error), + flow: "new", + threadId, + assistantMsgId: "no-persist-assistant", + workspaceId, + }); + return; + } + if ( + selection.filesystem_mode === "desktop_local_folder" && + (!selection.local_filesystem_mounts || selection.local_filesystem_mounts.length === 0) + ) { + toast.error("Select a local folder before using Local Folder mode."); + return; + } + + // Lazy thread creation: create thread on first message if it doesn't exist. + let currentThreadId = threadId; + let isNewThread = false; + if (!currentThreadId) { + try { + const newThread = await createThread(workspaceId, "New Chat"); + currentThreadId = newThread.id; + view.setThreadId(currentThreadId); + view.setCurrentThread(newThread); + queryClient.setQueryData(cacheKeys.threads.detail(newThread.id), newThread); + queryClient.setQueryData(cacheKeys.threads.messages(newThread.id), { messages: [] }); + + trackChatCreated(workspaceId, currentThreadId); + + isNewThread = true; + // Update URL silently using browser API (not router.replace) to avoid + // interrupting the ongoing fetch/streaming with React navigation. + window.history.replaceState( + null, + "", + `/dashboard/${workspaceId}/new-chat/${currentThreadId}` + ); + } catch (error) { + console.error("[chat-engine] Failed to create thread:", error); + await handleChatFailure({ + error: tagPreAcceptSendFailure(error), + flow: "new", + threadId: currentThreadId, + assistantMsgId: "no-persist-assistant", + workspaceId, + }); + return; + } + } + + // Seed the durable per-thread overlay with the pre-turn conversation and + // flip it to running so the page renders the live stream. + const streamThreadId = currentThreadId; + chatStreamStore.begin(streamThreadId, priorMessages); + + if (urlsSnapshot.length > 0) { + jotaiStore.set(pendingUserImageDataUrlsAtom, (prev) => + prev.filter((u) => !urlsSnapshot.includes(u)) + ); + } + + // Add user message to state. Mutable because the SSE + // ``data-user-message-id`` handler renames this optimistic id to the + // canonical ``msg-{db_id}`` once ``persist_user_turn`` resolves. + let userMsgId = `msg-user-${Date.now()}`; + + const authorMetadata = currentUser + ? { + custom: { + author: { + displayName: currentUser.display_name ?? null, + avatarUrl: currentUser.avatar_url ?? null, + }, + }, + } + : undefined; + + const existingImageUrls = new Set( + message.content + .filter( + (p): p is { type: "image"; image: string } => + typeof p === "object" && p !== null && "type" in p && p.type === "image" && "image" in p + ) + .map((p) => p.image) + ); + const extraImageParts = urlsSnapshot + .filter((u) => !existingImageUrls.has(u)) + .map((image) => ({ type: "image" as const, image })); + const userDisplayContent = [...message.content, ...extraImageParts]; + + const userMessage: ThreadMessageLike = { + id: userMsgId, + role: "user", + content: userDisplayContent, + createdAt: new Date(), + metadata: authorMetadata, + }; + chatStreamStore.setMessages(streamThreadId, (prev) => [...prev, userMessage]); + + trackChatMessageSent(workspaceId, streamThreadId, { + hasAttachments: userImages.length > 0, + hasMentionedDocuments: + mentionPayload.document_ids.length > 0 || + mentionPayload.folder_ids.length > 0 || + mentionPayload.connector_ids.length > 0, + messageLength: userQuery.length, + }); + + // Collect unique mention chips for display & persistence. + const allMentionedDocs: MentionedDocumentInfo[] = []; + const seenDocKeys = new Set(); + for (const doc of activeMentions) { + const key = getMentionDocKey(doc); + if (seenDocKeys.has(key)) continue; + seenDocKeys.add(key); + allMentionedDocs.push(doc); + } + + if (allMentionedDocs.length > 0) { + jotaiStore.set(messageDocumentsMapAtom, (prev) => ({ + ...prev, + [userMsgId]: allMentionedDocs, + })); + } + + const controller = new AbortController(); + chatStreamStore.beginActive(streamThreadId, controller); + + // Prepare assistant message. Mutable for the same reason as ``userMsgId``. + let assistantMsgId = `msg-assistant-${Date.now()}`; + const currentThinkingSteps = new Map(); + const contentPartsState: ContentPartsState = { + contentParts: [], + currentTextPartIndex: -1, + currentReasoningPartIndex: -1, + toolCallIndices: new Map(), + }; + const { contentParts } = contentPartsState; + let wasInterrupted = false; + let newAccepted = false; + let streamBatcher: FrameBatchedUpdater | null = null; + + try { + // Build message history for context. + const messageHistory = priorMessages + .filter((m) => m.role === "user" || m.role === "assistant") + .map((m) => { + let text = ""; + for (const part of m.content) { + if (typeof part === "object" && part.type === "text" && "text" in part) { + text += part.text; + } + } + return { role: m.role, content: text }; + }) + .filter((m) => m.content.length > 0); + + const hasDocumentIds = mentionPayload.document_ids.length > 0; + const hasFolderIds = mentionPayload.folder_ids.length > 0; + const hasConnectorIds = mentionPayload.connector_ids.length > 0; + const hasThreadIds = mentionPayload.thread_ids.length > 0; + + const response = await fetchWithTurnCancellingRetry(() => + authenticatedFetch(buildBackendUrl("/api/v1/new_chat"), { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + chat_id: streamThreadId, + user_query: userQuery.trim(), + workspace_id: workspaceId, + filesystem_mode: selection.filesystem_mode, + client_platform: selection.client_platform, + local_filesystem_mounts: selection.local_filesystem_mounts, + messages: messageHistory, + mentioned_document_ids: hasDocumentIds ? mentionPayload.document_ids : undefined, + mentioned_folder_ids: hasFolderIds ? mentionPayload.folder_ids : undefined, + mentioned_connector_ids: hasConnectorIds ? mentionPayload.connector_ids : undefined, + mentioned_connectors: hasConnectorIds ? mentionPayload.connectors : undefined, + mentioned_thread_ids: hasThreadIds ? mentionPayload.thread_ids : undefined, + mentioned_documents: allMentionedDocs.length > 0 ? allMentionedDocs : undefined, + disabled_tools: disabledTools.length > 0 ? disabledTools : undefined, + ...(userImages.length > 0 ? { user_images: userImages } : {}), + }), + signal: controller.signal, + }) + ); + + if (!response.ok) { + throw await toHttpResponseError(response); + } + newAccepted = true; + chatStreamStore.setMessages(streamThreadId, (prev) => [ + ...prev, + { + id: assistantMsgId, + role: "assistant", + content: [{ type: "text", text: "" }], + createdAt: new Date(), + }, + ]); + + const flushMessages = () => { + chatStreamStore.setMessages(streamThreadId, (prev) => + prev.map((m) => + m.id === assistantMsgId + ? { ...m, content: buildContentForUI(contentPartsState, toolsWithUI) } + : m + ) + ); + }; + const { batcher, scheduleFlush, forceFlush } = createStreamFlushHelpers(flushMessages); + streamBatcher = batcher; + + await consumeSseEvents(response, async (parsed) => { + if ( + processSharedStreamEvent(parsed, { + contentPartsState, + toolsWithUI, + currentThinkingSteps, + scheduleFlush, + forceFlush, + onTokenUsage: (data) => { + tokenUsageStore.set(assistantMsgId, data); + }, + onTurnStatus: (data) => { + if (data.status === "cancelling") { + chatStreamStore.markRecentCancel(); + } + }, + }) + ) { + return; + } + switch (parsed.type) { + case "data-thread-title-update": { + const titleData = parsed.data as { threadId: number; title: string }; + if (titleData?.title && titleData?.threadId === streamThreadId) { + view.setCurrentThread((prev) => (prev ? { ...prev, title: titleData.title } : prev)); + jotaiStore.set(updateChatTabTitleAtom, { + chatId: streamThreadId, + title: titleData.title, + }); + queryClient.setQueriesData( + { queryKey: ["threads", String(workspaceId)] }, + (old) => { + if (!old) return old; + const updateTitle = (list: ThreadListItem[]) => + list.map((t) => + t.id === titleData.threadId ? { ...t, title: titleData.title } : t + ); + return { + ...old, + threads: updateTitle(old.threads), + archived_threads: updateTitle(old.archived_threads), + }; + } + ); + } + break; + } + + case "data-documents-updated": { + const docEvent = parsed.data as { + action: string; + document: AgentCreatedDocument; + }; + if (docEvent?.document?.id) { + jotaiStore.set(agentCreatedDocumentsAtom, (prev) => { + if (prev.some((d) => d.id === docEvent.document.id)) return prev; + return [...prev, docEvent.document]; + }); + } + break; + } + + case "data-interrupt-request": { + wasInterrupted = true; + const interruptData = parsed.data as Record; + const actionRequests = (interruptData.action_requests ?? []) as Array<{ + name: string; + args: Record; + }>; + const paired = pairBundleToolCallIds( + contentPartsState.toolCallIndices, + contentPartsState.contentParts, + actionRequests + ); + const bundleToolCallIds: string[] = []; + for (let i = 0; i < actionRequests.length; i++) { + const action = actionRequests[i]; + let targetTcId = paired[i]; + if (!targetTcId) { + targetTcId = freshSynthToolCallId(contentPartsState.toolCallIndices, action.name, i); + addToolCall( + contentPartsState, + toolsWithUI, + targetTcId, + action.name, + action.args, + true + ); + } + updateToolCall(contentPartsState, targetTcId, { + result: { __interrupt__: true, ...interruptData }, + }); + bundleToolCallIds.push(targetTcId); + } + chatStreamStore.setMessages(streamThreadId, (prev) => + prev.map((m) => + m.id === assistantMsgId + ? { ...m, content: buildContentForUI(contentPartsState, toolsWithUI) } + : m + ) + ); + // ``tool_call_id`` is stamped on the backend by + // ``checkpointed_subagent_middleware``. Without it we can't + // address the paused subagent on resume — skip rather than + // fabricate a synthetic key. + const interruptId = String(interruptData.tool_call_id ?? ""); + if (interruptId) { + const incoming: PendingInterruptState = { + interruptId, + threadId: streamThreadId, + assistantMsgId, + interruptData, + bundleToolCallIds, + }; + chatStreamStore.setPendingInterrupts(streamThreadId, (prev) => { + const without = prev.filter((p) => p.interruptId !== interruptId); + return [...without, incoming]; + }); + } + break; + } + + case "data-action-log": { + applyActionLogSse(queryClient, streamThreadId, workspaceId, parsed.data); + break; + } + + case "data-action-log-updated": { + applyActionLogUpdatedSse( + queryClient, + streamThreadId, + parsed.data.id, + parsed.data.reversible + ); + break; + } + + case "data-turn-info": { + const turnId = readStreamedChatTurnId(parsed.data); + if (turnId) { + chatStreamStore.setMessages(streamThreadId, (prev) => + applyTurnIdToAssistantMessageList(prev, assistantMsgId, turnId) + ); + } + break; + } + + case "data-user-message-id": { + const parsedMsg = readStreamedMessageId(parsed.data); + if (!parsedMsg) break; + const newUserMsgId = `msg-${parsedMsg.messageId}`; + const oldUserMsgId = userMsgId; + chatStreamStore.setMessages(streamThreadId, (prev) => + prev.map((m) => + m.id === oldUserMsgId + ? mergeChatTurnIdIntoMessage({ ...m, id: newUserMsgId }, parsedMsg.turnId) + : m + ) + ); + if (allMentionedDocs.length > 0) { + jotaiStore.set(messageDocumentsMapAtom, (prev) => { + if (!(oldUserMsgId in prev)) { + return { ...prev, [newUserMsgId]: allMentionedDocs }; + } + const { [oldUserMsgId]: _removed, ...rest } = prev; + return { ...rest, [newUserMsgId]: allMentionedDocs }; + }); + } + userMsgId = newUserMsgId; + if (isNewThread) { + queryClient.invalidateQueries({ + queryKey: ["threads", String(workspaceId)], + }); + } + break; + } + + case "data-assistant-message-id": { + const parsedMsg = readStreamedMessageId(parsed.data); + if (!parsedMsg) break; + const newAssistantMsgId = `msg-${parsedMsg.messageId}`; + const oldAssistantMsgId = assistantMsgId; + tokenUsageStore.rename(oldAssistantMsgId, newAssistantMsgId); + chatStreamStore.setMessages(streamThreadId, (prev) => + prev.map((m) => + m.id === oldAssistantMsgId + ? mergeChatTurnIdIntoMessage({ ...m, id: newAssistantMsgId }, parsedMsg.turnId) + : m + ) + ); + chatStreamStore.setPendingInterrupts(streamThreadId, (prev) => + prev.map((p) => + p.assistantMsgId === oldAssistantMsgId + ? { ...p, assistantMsgId: newAssistantMsgId } + : p + ) + ); + assistantMsgId = newAssistantMsgId; + break; + } + } + }); + + batcher.flush(); + + if (contentParts.length > 0 && !wasInterrupted) { + trackChatResponseReceived(workspaceId, streamThreadId); + } + } catch (error) { + streamBatcher?.dispose(); + await handleStreamTerminalError({ + error, + flow: "new", + threadId: streamThreadId, + assistantMsgId, + accepted: newAccepted, + workspaceId, + onPreAcceptFailure: async () => { + // Pre-accept failure means the BE never accepted the request — no + // server-side persistence ran. Roll back the optimistic UI. + chatStreamStore.setMessages(streamThreadId, (prev) => + prev.filter((m) => m.id !== userMsgId) + ); + jotaiStore.set(messageDocumentsMapAtom, (prev) => { + if (!(userMsgId in prev)) return prev; + const { [userMsgId]: _removed, ...rest } = prev; + return rest; + }); + }, + }); + } finally { + chatStreamStore.setRunning(streamThreadId, false); + chatStreamStore.clearActive(controller); + void queryClient.invalidateQueries({ + queryKey: cacheKeys.threads.messages(streamThreadId), + }); + void queryClient.invalidateQueries({ + queryKey: cacheKeys.threads.detail(streamThreadId), + }); + } +} + +// --------------------------------------------------------------------------- +// Resume (HITL decisions) +// --------------------------------------------------------------------------- + +export async function resumeChat( + ctx: EngineContext, + decisions: Array<{ + type: string; + message?: string; + edited_action?: { name: string; args: Record }; + }> +): Promise { + const { workspaceId, threadId } = ctx; + if (threadId == null) return; + const pendingInterrupts = chatStreamStore.getPendingInterrupts(threadId); + if (pendingInterrupts.length === 0) return; + + const resumeThreadId = pendingInterrupts[0].threadId; + let assistantMsgId = pendingInterrupts[0].assistantMsgId; + const allBundleToolCallIds = pendingInterrupts.flatMap((p) => p.bundleToolCallIds); + chatStreamStore.setPendingInterrupts(resumeThreadId, () => []); + chatStreamStore.setRunning(resumeThreadId, true); + + const controller = new AbortController(); + chatStreamStore.beginActive(resumeThreadId, controller); + + const localFilesystemEnabled = + jotaiStore.get(agentFlagsAtom).data?.enable_desktop_local_filesystem === true; + const disabledTools = jotaiStore.get(disabledToolsAtom); + + const currentThinkingSteps = new Map(); + const contentPartsState: ContentPartsState = { + contentParts: [], + currentTextPartIndex: -1, + currentReasoningPartIndex: -1, + toolCallIndices: new Map(), + }; + const { contentParts, toolCallIndices } = contentPartsState; + let resumeAccepted = false; + let streamBatcher: FrameBatchedUpdater | null = null; + + const existingMsg = chatStreamStore + .getMessages(resumeThreadId) + .find((m) => m.id === assistantMsgId); + if (existingMsg && Array.isArray(existingMsg.content)) { + contentPartsState.suppressStepSeparators = true; + for (const part of existingMsg.content) { + if (typeof part === "object" && part !== null) { + const p = part as Record; + if (p.type === "text") { + contentParts.push({ type: "text", text: String(p.text ?? "") }); + contentPartsState.currentTextPartIndex = contentParts.length - 1; + } else if (p.type === "tool-call") { + toolCallIndices.set(String(p.toolCallId), contentParts.length); + contentParts.push({ + type: "tool-call", + toolCallId: String(p.toolCallId), + toolName: String(p.toolName), + args: (p.args as Record) ?? {}, + result: p.result as unknown, + ...(typeof p.argsText === "string" ? { argsText: p.argsText } : {}), + ...(typeof p.langchainToolCallId === "string" + ? { langchainToolCallId: p.langchainToolCallId } + : {}), + ...(p.metadata && typeof p.metadata === "object" + ? { metadata: p.metadata as Record } + : {}), + }); + contentPartsState.currentTextPartIndex = -1; + } else if (p.type === "data-thinking-steps") { + const stepsData = p.data as { steps: ThinkingStepData[] } | undefined; + contentParts.push({ + type: "data-thinking-steps", + data: { steps: stepsData?.steps ?? [] }, + }); + for (const step of stepsData?.steps ?? []) { + currentThinkingSteps.set(step.id, step); + } + } + } + } + } + + // Apply each decision to its own card by toolCallId so mixed bundles + // (approve/edit/reject) do not collapse onto ``decisions[0]``. + const decisionByTcId = new Map(); + const tcIds = allBundleToolCallIds; + if (decisions.length === tcIds.length) { + for (let i = 0; i < tcIds.length; i++) decisionByTcId.set(tcIds[i], decisions[i]); + } + if (decisionByTcId.size > 0) { + for (const part of contentParts) { + if (part.type !== "tool-call") continue; + const tcId = part.toolCallId as string | undefined; + const d = tcId ? decisionByTcId.get(tcId) : undefined; + if (!d) continue; + if (typeof part.result !== "object" || part.result === null) continue; + if (!("__interrupt__" in (part.result as Record))) continue; + const decided = d.type; + if (decided === "edit" && d.edited_action) { + const mergedArgs = { ...part.args, ...d.edited_action.args }; + part.args = mergedArgs; + part.argsText = JSON.stringify(mergedArgs, null, 2); + } + part.result = { + ...(part.result as Record), + __decided__: decided, + }; + } + } + + try { + const selection = await getAgentFilesystemSelection(workspaceId, { localFilesystemEnabled }); + const response = await fetchWithTurnCancellingRetry(() => + authenticatedFetch(buildBackendUrl(`/api/v1/threads/${resumeThreadId}/resume`), { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + workspace_id: workspaceId, + decisions, + disabled_tools: disabledTools.length > 0 ? disabledTools : undefined, + filesystem_mode: selection.filesystem_mode, + client_platform: selection.client_platform, + local_filesystem_mounts: selection.local_filesystem_mounts, + }), + signal: controller.signal, + }) + ); + + if (!response.ok) { + throw await toHttpResponseError(response); + } + resumeAccepted = true; + + const flushMessages = () => { + chatStreamStore.setMessages(resumeThreadId, (prev) => + prev.map((m) => + m.id === assistantMsgId + ? { ...m, content: buildContentForUI(contentPartsState, toolsWithUI) } + : m + ) + ); + }; + const { batcher, scheduleFlush, forceFlush } = createStreamFlushHelpers(flushMessages); + streamBatcher = batcher; + + await consumeSseEvents(response, async (parsed) => { + if ( + processSharedStreamEvent(parsed, { + contentPartsState, + toolsWithUI, + currentThinkingSteps, + scheduleFlush, + forceFlush, + onTokenUsage: (data) => { + tokenUsageStore.set(assistantMsgId, data); + }, + onTurnStatus: (data) => { + if (data.status === "cancelling") { + chatStreamStore.markRecentCancel(); + } + }, + }) + ) { + return; + } + switch (parsed.type) { + case "data-interrupt-request": { + const interruptData = parsed.data as Record; + const actionRequests = (interruptData.action_requests ?? []) as Array<{ + name: string; + args: Record; + }>; + const paired = pairBundleToolCallIds( + contentPartsState.toolCallIndices, + contentPartsState.contentParts, + actionRequests + ); + const bundleToolCallIds: string[] = []; + for (let i = 0; i < actionRequests.length; i++) { + const action = actionRequests[i]; + let targetTcId = paired[i]; + if (!targetTcId) { + targetTcId = freshSynthToolCallId(contentPartsState.toolCallIndices, action.name, i); + addToolCall( + contentPartsState, + toolsWithUI, + targetTcId, + action.name, + action.args, + true + ); + } + updateToolCall(contentPartsState, targetTcId, { + result: { __interrupt__: true, ...interruptData }, + }); + bundleToolCallIds.push(targetTcId); + } + chatStreamStore.setMessages(resumeThreadId, (prev) => + prev.map((m) => + m.id === assistantMsgId + ? { ...m, content: buildContentForUI(contentPartsState, toolsWithUI) } + : m + ) + ); + { + const interruptId = String(interruptData.tool_call_id ?? ""); + if (interruptId) { + const incoming: PendingInterruptState = { + interruptId, + threadId: resumeThreadId, + assistantMsgId, + interruptData, + bundleToolCallIds, + }; + chatStreamStore.setPendingInterrupts(resumeThreadId, (prev) => { + const without = prev.filter((p) => p.interruptId !== interruptId); + return [...without, incoming]; + }); + } + } + break; + } + + case "data-action-log": { + applyActionLogSse(queryClient, resumeThreadId, workspaceId, parsed.data); + break; + } + + case "data-action-log-updated": { + applyActionLogUpdatedSse( + queryClient, + resumeThreadId, + parsed.data.id, + parsed.data.reversible + ); + break; + } + + case "data-turn-info": { + const turnId = readStreamedChatTurnId(parsed.data); + if (turnId) { + chatStreamStore.setMessages(resumeThreadId, (prev) => + applyTurnIdToAssistantMessageList(prev, assistantMsgId, turnId) + ); + } + break; + } + + case "data-assistant-message-id": { + const parsedMsg = readStreamedMessageId(parsed.data); + if (!parsedMsg) break; + const newAssistantMsgId = `msg-${parsedMsg.messageId}`; + const oldAssistantMsgId = assistantMsgId; + tokenUsageStore.rename(oldAssistantMsgId, newAssistantMsgId); + chatStreamStore.setMessages(resumeThreadId, (prev) => + prev.map((m) => + m.id === oldAssistantMsgId + ? mergeChatTurnIdIntoMessage({ ...m, id: newAssistantMsgId }, parsedMsg.turnId) + : m + ) + ); + assistantMsgId = newAssistantMsgId; + break; + } + } + }); + + batcher.flush(); + } catch (error) { + streamBatcher?.dispose(); + await handleStreamTerminalError({ + error, + flow: "resume", + threadId: resumeThreadId, + assistantMsgId, + accepted: resumeAccepted, + workspaceId, + }); + } finally { + chatStreamStore.setRunning(resumeThreadId, false); + chatStreamStore.clearActive(controller); + void queryClient.invalidateQueries({ + queryKey: cacheKeys.threads.messages(resumeThreadId), + }); + void queryClient.invalidateQueries({ + queryKey: cacheKeys.threads.detail(resumeThreadId), + }); + } +} + +// --------------------------------------------------------------------------- +// Regenerate (edit / reload) +// --------------------------------------------------------------------------- + +export async function regenerateChat( + ctx: EngineContext, + newUserQuery: string | null, + editExtras?: { + userMessageContent: ThreadMessageLike["content"]; + userImages: NewChatUserImagePayload[]; + sourceUserMessageId?: string; + }, + editFromPosition?: { + fromMessageId?: number | null; + revertActions?: boolean; + } +): Promise { + const { workspaceId, threadId, priorMessages } = ctx; + if (!threadId) { + toast.error("Cannot regenerate: no active chat thread"); + return; + } + const streamThreadId = threadId; + + const isEdit = newUserQuery !== null; + + // Supersede any previous in-flight turn. + chatStreamStore.abortActive(); + + const messageDocumentsMap = jotaiStore.get(messageDocumentsMapAtom); + const localFilesystemEnabled = + jotaiStore.get(agentFlagsAtom).data?.enable_desktop_local_filesystem === true; + const disabledTools = jotaiStore.get(disabledToolsAtom); + + // Extract the original user query BEFORE removing messages (reload mode). + let userQueryToDisplay: string | undefined; + let originalUserMessageContent: ThreadMessageLike["content"] | null = null; + let originalUserMessageMetadata: ThreadMessageLike["metadata"] | undefined; + let sourceUserMessageId: string | undefined = editExtras?.sourceUserMessageId; + + if (!isEdit) { + const lastUserMessage = [...priorMessages].reverse().find((m) => m.role === "user"); + if (lastUserMessage) { + sourceUserMessageId = lastUserMessage.id; + originalUserMessageContent = lastUserMessage.content; + originalUserMessageMetadata = lastUserMessage.metadata; + for (const part of lastUserMessage.content) { + if (typeof part === "object" && part.type === "text" && "text" in part) { + userQueryToDisplay = part.text; + break; + } + } + } + } else { + userQueryToDisplay = newUserQuery; + } + + // Seed the durable overlay with the pre-regenerate conversation and flip + // to running so the page renders the live stream. + chatStreamStore.begin(streamThreadId, priorMessages); + + const controller = new AbortController(); + chatStreamStore.beginActive(streamThreadId, controller); + + let userMsgId = `msg-user-${Date.now()}`; + let assistantMsgId = `msg-assistant-${Date.now()}`; + const currentThinkingSteps = new Map(); + + const contentPartsState: ContentPartsState = { + contentParts: [], + currentTextPartIndex: -1, + currentReasoningPartIndex: -1, + toolCallIndices: new Map(), + }; + const { contentParts } = contentPartsState; + let regenerateAccepted = false; + let streamBatcher: FrameBatchedUpdater | null = null; + + const userMessage: ThreadMessageLike = { + id: userMsgId, + role: "user", + content: isEdit + ? (editExtras?.userMessageContent ?? [{ type: "text", text: newUserQuery ?? "" }]) + : originalUserMessageContent || [{ type: "text", text: userQueryToDisplay || "" }], + createdAt: new Date(), + metadata: isEdit ? undefined : originalUserMessageMetadata, + }; + const sourceMentionedDocs = + sourceUserMessageId && messageDocumentsMap[sourceUserMessageId] + ? messageDocumentsMap[sourceUserMessageId] + : []; + try { + const selection = await getAgentFilesystemSelection(workspaceId, { localFilesystemEnabled }); + const regenerateDocIds = sourceMentionedDocs.filter((d) => d.kind === "doc").map((d) => d.id); + const regenerateFolderIds = sourceMentionedDocs + .filter((d) => d.kind === "folder") + .map((d) => d.id); + const regenerateConnectors = sourceMentionedDocs.filter((d) => d.kind === "connector"); + const regenerateThreadIds = sourceMentionedDocs + .filter((d) => d.kind === "thread") + .map((d) => d.id); + + const requestBody: Record = { + workspace_id: workspaceId, + 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, + mentioned_document_ids: regenerateDocIds.length > 0 ? regenerateDocIds : undefined, + mentioned_folder_ids: regenerateFolderIds.length > 0 ? regenerateFolderIds : undefined, + mentioned_connector_ids: + regenerateConnectors.length > 0 ? regenerateConnectors.map((d) => d.id) : undefined, + mentioned_connectors: regenerateConnectors.length > 0 ? regenerateConnectors : undefined, + mentioned_thread_ids: regenerateThreadIds.length > 0 ? regenerateThreadIds : undefined, + mentioned_documents: sourceMentionedDocs.length > 0 ? sourceMentionedDocs : undefined, + }; + if (isEdit) { + requestBody.user_images = editExtras?.userImages ?? []; + } + if (editFromPosition?.fromMessageId != null) { + requestBody.from_message_id = editFromPosition.fromMessageId; + if (editFromPosition.revertActions) { + requestBody.revert_actions = true; + } + } + const response = await fetchWithTurnCancellingRetry(() => + authenticatedFetch(getRegenerateUrl(streamThreadId), { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(requestBody), + signal: controller.signal, + }) + ); + + if (!response.ok) { + throw await toHttpResponseError(response); + } + regenerateAccepted = true; + + chatStreamStore.setMessages(streamThreadId, (prev) => { + let base = prev; + if (editFromPosition?.fromMessageId != null) { + const targetId = `msg-${editFromPosition.fromMessageId}`; + const sliceIndex = prev.findIndex((m) => m.id === targetId); + if (sliceIndex >= 0) { + base = prev.slice(0, sliceIndex); + } + } else if (prev.length >= 2) { + base = prev.slice(0, -2); + } + return [ + ...base, + userMessage, + { + id: assistantMsgId, + role: "assistant", + content: [{ type: "text", text: "" }], + createdAt: new Date(), + }, + ]; + }); + if (sourceMentionedDocs.length > 0) { + jotaiStore.set(messageDocumentsMapAtom, (prev) => ({ + ...prev, + [userMsgId]: sourceMentionedDocs, + })); + } + + const flushMessages = () => { + chatStreamStore.setMessages(streamThreadId, (prev) => + prev.map((m) => + m.id === assistantMsgId + ? { ...m, content: buildContentForUI(contentPartsState, toolsWithUI) } + : m + ) + ); + }; + const { batcher, scheduleFlush, forceFlush } = createStreamFlushHelpers(flushMessages); + streamBatcher = batcher; + + await consumeSseEvents(response, async (parsed) => { + if ( + processSharedStreamEvent(parsed, { + contentPartsState, + toolsWithUI, + currentThinkingSteps, + scheduleFlush, + forceFlush, + onTokenUsage: (data) => { + tokenUsageStore.set(assistantMsgId, data); + }, + onTurnStatus: (data) => { + if (data.status === "cancelling") { + chatStreamStore.markRecentCancel(); + } + }, + }) + ) { + return; + } + switch (parsed.type) { + case "data-action-log": { + applyActionLogSse(queryClient, streamThreadId, workspaceId, parsed.data); + break; + } + + case "data-action-log-updated": { + applyActionLogUpdatedSse( + queryClient, + streamThreadId, + parsed.data.id, + parsed.data.reversible + ); + break; + } + + case "data-turn-info": { + const turnId = readStreamedChatTurnId(parsed.data); + if (turnId) { + chatStreamStore.setMessages(streamThreadId, (prev) => + applyTurnIdToAssistantMessageList(prev, assistantMsgId, turnId) + ); + } + break; + } + + case "data-user-message-id": { + const parsedMsg = readStreamedMessageId(parsed.data); + if (!parsedMsg) break; + const newUserMsgId = `msg-${parsedMsg.messageId}`; + const oldUserMsgId = userMsgId; + chatStreamStore.setMessages(streamThreadId, (prev) => + prev.map((m) => + m.id === oldUserMsgId + ? mergeChatTurnIdIntoMessage({ ...m, id: newUserMsgId }, parsedMsg.turnId) + : m + ) + ); + if (sourceMentionedDocs.length > 0) { + jotaiStore.set(messageDocumentsMapAtom, (prev) => { + if (!(oldUserMsgId in prev)) { + return { ...prev, [newUserMsgId]: sourceMentionedDocs }; + } + const { [oldUserMsgId]: _removed, ...rest } = prev; + return { ...rest, [newUserMsgId]: sourceMentionedDocs }; + }); + } + userMsgId = newUserMsgId; + break; + } + + case "data-assistant-message-id": { + const parsedMsg = readStreamedMessageId(parsed.data); + if (!parsedMsg) break; + const newAssistantMsgId = `msg-${parsedMsg.messageId}`; + const oldAssistantMsgId = assistantMsgId; + tokenUsageStore.rename(oldAssistantMsgId, newAssistantMsgId); + chatStreamStore.setMessages(streamThreadId, (prev) => + prev.map((m) => + m.id === oldAssistantMsgId + ? mergeChatTurnIdIntoMessage({ ...m, id: newAssistantMsgId }, parsedMsg.turnId) + : m + ) + ); + assistantMsgId = newAssistantMsgId; + break; + } + + case "data-revert-results": { + const summary = parsed.data; + const failureCount = + summary.failed + summary.not_reversible + (summary.permission_denied ?? 0); + if (failureCount > 0) { + toast.warning( + `Pre-revert: ${summary.reverted}/${summary.total} undone, ${failureCount} could not be rolled back.` + ); + } else if (summary.reverted > 0) { + toast.success( + summary.reverted === 1 + ? "Reverted 1 downstream action before regenerating." + : `Reverted ${summary.reverted} downstream actions before regenerating.` + ); + } + for (const r of summary.results) { + if (r.status === "reverted" || r.status === "already_reverted") { + markActionRevertedInCache( + queryClient, + streamThreadId, + r.action_id, + r.new_action_id ?? null + ); + } + } + break; + } + } + }); + + batcher.flush(); + + if (contentParts.length > 0) { + trackChatResponseReceived(workspaceId, streamThreadId); + } + } catch (error) { + streamBatcher?.dispose(); + await handleStreamTerminalError({ + error, + flow: "regenerate", + threadId: streamThreadId, + assistantMsgId, + accepted: regenerateAccepted, + workspaceId, + }); + } finally { + chatStreamStore.setRunning(streamThreadId, false); + chatStreamStore.clearActive(controller); + void queryClient.invalidateQueries({ + queryKey: cacheKeys.threads.messages(streamThreadId), + }); + void queryClient.invalidateQueries({ + queryKey: cacheKeys.threads.detail(streamThreadId), + }); + } +} + +export type { HitlDecision };