diff --git a/surfsense_web/app/dashboard/[workspace_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[workspace_id]/new-chat/[[...chat_id]]/page.tsx index 4bb355049..f5b857ce1 100644 --- a/surfsense_web/app/dashboard/[workspace_id]/new-chat/[[...chat_id]]/page.tsx +++ b/surfsense_web/app/dashboard/[workspace_id]/new-chat/[[...chat_id]]/page.tsx @@ -7,14 +7,11 @@ import { useExternalStoreRuntime, } from "@assistant-ui/react"; import { useQueryClient } from "@tanstack/react-query"; -import { useAtomValue, useSetAtom, useStore } from "jotai"; +import { useAtomValue, useSetAtom } from "jotai"; import dynamic from "next/dynamic"; import { useParams } from "next/navigation"; import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { toast } from "sonner"; -import { z } from "zod"; -import { agentFlagsAtom } from "@/atoms/agent/agent-flags-query.atom"; -import { disabledToolsAtom } from "@/atoms/agent-tools/agent-tools.atoms"; import { clearTargetCommentIdAtom, currentThreadAtom, @@ -22,24 +19,15 @@ import { setTargetCommentIdAtom, } from "@/atoms/chat/current-thread.atom"; import { - deriveMentionedPayload, type MentionedDocumentInfo, mentionedDocumentsAtom, messageDocumentsMapAtom, - submittedMentionsAtom, } from "@/atoms/chat/mentioned-documents.atom"; -import { pendingUserImageDataUrlsAtom } from "@/atoms/chat/pending-user-images.atom"; -import { - clearPlanOwnerRegistry, - // extractWriteTodosFromContent, -} from "@/atoms/chat/plan-state.atom"; -import { setPremiumAlertForThreadAtom } from "@/atoms/chat/premium-alert.atom"; +import { clearPlanOwnerRegistry } from "@/atoms/chat/plan-state.atom"; import { closeReportPanelAtom } from "@/atoms/chat/report-panel.atom"; -import { type AgentCreatedDocument, agentCreatedDocumentsAtom } from "@/atoms/documents/ui.atoms"; import { closeEditorPanelAtom } from "@/atoms/editor/editor-panel.atom"; import { membersAtom } from "@/atoms/members/members-query.atoms"; -import { removeChatTabAtom, syncChatTabAtom, updateChatTabTitleAtom } from "@/atoms/tabs/tabs.atom"; -import { currentUserAtom } from "@/atoms/user/user-query.atoms"; +import { removeChatTabAtom, syncChatTabAtom } from "@/atoms/tabs/tabs.atom"; import { EditMessageDialog, type EditMessageDialogChoice, @@ -47,7 +35,6 @@ import { import { StepSeparatorDataUI } from "@/components/assistant-ui/step-separator"; import { Thread } from "@/components/assistant-ui/thread"; import { - createTokenUsageStore, type TokenUsageData, TokenUsageProvider, } from "@/components/assistant-ui/token-usage-context"; @@ -60,64 +47,31 @@ import { type PendingInterruptState, } from "@/features/chat-messages/hitl"; import { TimelineDataUI } from "@/features/chat-messages/timeline"; -import { - applyActionLogSse, - applyActionLogUpdatedSse, - markActionRevertedInCache, - useAgentActionsQuery, -} from "@/hooks/use-agent-actions-query"; +import { useAgentActionsQuery } from "@/hooks/use-agent-actions-query"; import { useChatSessionStateSync } from "@/hooks/use-chat-session-state"; import { useMessagesSync } from "@/hooks/use-messages-sync"; import { useThreadDetail, useThreadMessages } from "@/hooks/use-thread-queries"; -import { getAgentFilesystemSelection } from "@/lib/agent-filesystem"; import { documentsApiService } from "@/lib/apis/documents-api.service"; -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 { convertToThreadMessage, reconcileInterruptedAssistantMessages, } from "@/lib/chat/message-utils"; -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, - type ToolUIGate, - updateToolCall, -} from "@/lib/chat/streaming-state"; -import { - appendMessage, - createThread, - getRegenerateUrl, - type ThreadListItem, - type ThreadListResponse, - type ThreadRecord, -} from "@/lib/chat/thread-persistence"; + cancelActiveTurn, + type EngineContext, + regenerateChat, + resumeChat, + startNewChat, +} from "@/lib/chat/stream-engine/engine"; +import { extractMentionedDocuments } from "@/lib/chat/stream-engine/helpers"; +import { chatStreamStore } from "@/lib/chat/stream-engine/store"; +import { useChatStream } from "@/lib/chat/stream-engine/use-chat-stream"; +import 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 { NotFoundError } from "@/lib/error"; -import { - trackChatBlocked, - trackChatCreated, - trackChatErrorDetailed, - trackChatMessageSent, - trackChatResponseReceived, -} from "@/lib/posthog/events"; -import { cacheKeys } from "@/lib/query-client/cache-keys"; const MobileEditorPanel = dynamic( () => @@ -148,156 +102,8 @@ const MobileArtifactsPanel = dynamic( { ssr: false } ); -/** - * Generate a synthetic ``toolCallId`` for an action_request that has no - * matching streamed tool-call card (HITL-blocked subagent calls don't surface - * as tool-call events). Suffixes a counter when the base id is already taken - * — sequential interrupts for the same tool name otherwise collide on - * ``interrupt-${name}-${i}`` and crash assistant-ui with a duplicate-key error. - */ -function freshSynthToolCallId( - toolCallIndices: Map, - toolName: string, - index: number -): string { - const base = `interrupt-${toolName}-${index}`; - if (!toolCallIndices.has(base)) return base; - let n = 1; - while (toolCallIndices.has(`${base}-${n}`)) n++; - return `${base}-${n}`; -} - -/** - * Pair each ``action_request`` to a unique pending tool-call card, preserving - * order so ``decisions[i]`` lines up with ``action_requests[i]`` on the wire. - * - * Same-name bundles (e.g. three ``create_jira_issue``) used to collapse onto - * one card because the matcher keyed by name; this consumes each card via the - * ``claimed`` set and walks forward in DOM order. - */ -function pairBundleToolCallIds( - toolCallIndices: Map, - contentParts: Array<{ - type: string; - toolName?: string; - result?: unknown; - }>, - actionRequests: ReadonlyArray<{ name: string }> -): Array { - const claimed = new Set(); - const paired: Array = []; - for (const action of actionRequests) { - let matched: string | null = null; - for (const [tcId, idx] of toolCallIndices) { - if (claimed.has(tcId)) continue; - const part = contentParts[idx]; - if (!part || part.type !== "tool-call" || part.toolName !== action.name) continue; - const result = part.result as Record | undefined | null; - if (result == null || (result.__interrupt__ === true && !result.__decided__)) { - matched = tcId; - claimed.add(tcId); - break; - } - } - paired.push(matched); - } - return paired; -} - -/** - * Zod schema for mentioned document info (for type-safe parsing). - * - * ``kind`` defaults to ``"doc"`` so messages persisted before folder - * mentions existed deserialise unchanged. - */ -const MentionedDocumentInfoSchema = z.object({ - id: z.number(), - title: z.string(), - document_type: z.string().optional(), - kind: z - .union([z.literal("doc"), z.literal("folder"), z.literal("connector"), z.literal("thread")]) - .optional() - .default("doc"), - connector_type: z.string().optional(), - account_name: z.string().optional(), -}); - -const MentionedDocumentsPartSchema = z.object({ - type: z.literal("mentioned-documents"), - documents: z.array(MentionedDocumentInfoSchema), -}); - -/** - * Extract mentioned documents from message content (type-safe with Zod) - */ -function extractMentionedDocuments(content: unknown): MentionedDocumentInfo[] { - if (!Array.isArray(content)) return []; - - for (const part of content) { - const result = MentionedDocumentsPartSchema.safeParse(part); - if (result.success) { - return result.data.documents.map((doc) => { - if (doc.kind === "connector") { - return { - id: doc.id, - title: doc.title, - kind: "connector", - connector_type: doc.connector_type ?? doc.document_type ?? "UNKNOWN", - account_name: doc.account_name ?? doc.title, - }; - } - if (doc.kind === "folder") { - return { - id: doc.id, - title: doc.title, - kind: "folder", - }; - } - if (doc.kind === "thread") { - return { - id: doc.id, - title: doc.title, - kind: "thread", - }; - } - return { - id: doc.id, - title: doc.title, - document_type: doc.document_type ?? "UNKNOWN", - kind: "doc", - }; - }); - } - } - - return []; -} - -/** - * Every tool call renders a card. The legacy - * ``BASE_TOOLS_WITH_UI`` allowlist used to drop unknown tool calls on the - * floor; we now route everything through ``ToolFallback``. Persisted - * payload size stays bounded because the backend's - * ``format_thinking_step`` summarisation and the - * ``result_length``-only default for unknown tools (see - * ``stream_new_chat.py``) keep the JSON from ballooning. - */ -const TOOLS_WITH_UI_ALL: ToolUIGate = "all"; -const TURN_CANCELLING_INITIAL_DELAY_MS = 200; -const TURN_CANCELLING_BACKOFF_FACTOR = 2; -const TURN_CANCELLING_MAX_DELAY_MS = 1500; -const RECENT_CANCEL_WINDOW_MS = 5_000; - -function sleep(ms: number): Promise { - return new Promise((resolve) => setTimeout(resolve, ms)); -} - -function computeFallbackTurnCancellingRetryDelay(attempt: number): number { - const safeAttempt = Math.max(1, attempt); - const raw = - TURN_CANCELLING_INITIAL_DELAY_MS * TURN_CANCELLING_BACKOFF_FACTOR ** (safeAttempt - 1); - return Math.min(raw, TURN_CANCELLING_MAX_DELAY_MS); -} +/** Stable empty reference so idle threads don't re-render the interrupt provider. */ +const EMPTY_PENDING_INTERRUPTS: PendingInterruptState[] = []; function parseUrlChatId(id: string | string[] | undefined): number { let parsed = 0; @@ -372,103 +178,37 @@ export default function NewChatPage() { const activeThreadId = urlChatId > 0 ? urlChatId : threadId; const handledLoadErrorThreadRef = useRef(null); const [currentThread, setCurrentThread] = useState(null); + // DB-hydrated messages for the viewed thread (idle display). While a turn + // is streaming, the live overlay in ``chatStreamStore`` takes precedence + // (see ``displayMessages``) so it survives this page unmounting on nav. const [messages, setMessages] = useState([]); - const [isRunning, setIsRunning] = useState(false); - const [tokenUsageStore] = useState(() => createTokenUsageStore()); - const abortControllerRef = useRef(null); - const recentCancelRequestedAtRef = useRef(0); - // One entry per paused subagent, in receipt order (which matches the - // backend's ``state.interrupts`` traversal — and therefore the order - // ``slice_decisions_by_tool_call`` consumes on resume). Cleared on submit - // or on a fresh user turn. - const [pendingInterrupts, setPendingInterrupts] = useState([]); - // Per-card staged decisions held until every pending card has submitted, - // at which point we batch them into one ``hitl-decision`` event in the - // same order as ``pendingInterrupts``. Using a ref because partial - // progress should not re-render the page. - const stagedDecisionsByInterruptIdRef = useRef>(new Map()); - const toolsWithUI = TOOLS_WITH_UI_ALL; + + // Durable, cross-navigation streaming state for the viewed thread. + const streamState = useChatStream(activeThreadId); + const isRunning = streamState?.isRunning ?? false; + const pendingInterrupts = streamState?.pendingInterrupts ?? EMPTY_PENDING_INTERRUPTS; + // Live overlay while a turn is streaming / awaiting HITL; DB-hydrated + // messages once the overlay is cleared (the hydration effect drops it only + // after the DB catches up, so there is no finish->refetch gap). + const displayMessages = streamState ? streamState.messages : messages; + + // One shared token-usage store, alive across navigation. + const tokenUsageStore = chatStreamStore.tokenUsage; + const setMessageDocumentsMap = useSetAtom(messageDocumentsMapAtom); - - const persistAssistantErrorMessage = useCallback( - async ({ - threadId, - assistantMsgId, - text, - }: { - threadId: number | null; - assistantMsgId: string; - text: string; - }) => { - setMessages((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); - setMessages((prev) => - prev.map((m) => (m.id === assistantMsgId ? { ...m, id: newMsgId } : m)) - ); - } catch (persistErr) { - console.error("Failed to persist assistant error message:", persistErr); - } - }, - [tokenUsageStore] - ); - - // NOTE: ``persistUserTurn`` / ``persistAssistantTurn`` callbacks - // were removed in the SSE-based message ID handshake refactor. - // ``stream_new_chat`` and ``stream_resume_chat`` now persist both - // the user and assistant rows server-side via - // ``persist_user_turn`` / ``persist_assistant_shell`` and emit - // ``data-user-message-id`` / ``data-assistant-message-id`` SSE - // events; the consumers below rename the optimistic ids in real - // time. ``persistAssistantErrorMessage`` (above) is intentionally - // kept — it is the pre-stream-error fallback fired when the - // server NEVER accepted the request, and the BE has nothing to - // persist in that case. - - // Get disabled tools from the tool toggle UI - const disabledTools = useAtomValue(disabledToolsAtom); - - const jotaiStore = useStore(); - const mentionedDocuments = useAtomValue(mentionedDocumentsAtom); - const messageDocumentsMap = useAtomValue(messageDocumentsMapAtom); const setMentionedDocuments = useSetAtom(mentionedDocumentsAtom); const currentThreadState = useAtomValue(currentThreadAtom); const setCurrentThreadMetadata = useSetAtom(setCurrentThreadMetadataAtom); - const setPremiumAlertForThread = useSetAtom(setPremiumAlertForThreadAtom); const setTargetCommentId = useSetAtom(setTargetCommentIdAtom); const clearTargetCommentId = useSetAtom(clearTargetCommentIdAtom); const closeReportPanel = useSetAtom(closeReportPanelAtom); const closeEditorPanel = useSetAtom(closeEditorPanelAtom); const syncChatTab = useSetAtom(syncChatTabAtom); - const updateChatTabTitle = useSetAtom(updateChatTabTitleAtom); const removeChatTab = useSetAtom(removeChatTabAtom); - const setAgentCreatedDocuments = useSetAtom(agentCreatedDocumentsAtom); - const pendingUserImageUrls = useAtomValue(pendingUserImageDataUrlsAtom); - const setPendingUserImageUrls = useSetAtom(pendingUserImageDataUrlsAtom); - // Edit dialog state. Holds the message id being edited and - // the (already extracted) regenerate args so we can resume the edit - // after the user picks "revert all" / "continue" / "cancel". + + // Edit dialog state. Holds the message id being edited and the (already + // extracted) regenerate args so we can resume the edit after the user picks + // "revert all" / "continue" / "cancel". const [editDialogState, setEditDialogState] = useState<{ fromMessageId: number; userQuery: string | null; @@ -478,14 +218,17 @@ export default function NewChatPage() { downstreamTotalCount: number; } | null>(null); - // Get current user for author info in shared chats - const { data: currentUser } = useAtomValue(currentUserAtom); - const { data: agentFlags } = useAtomValue(agentFlagsAtom); - const localFilesystemEnabled = agentFlags?.enable_desktop_local_filesystem === true; + // Per-card staged decisions held until every pending card has submitted, at + // which point we batch them into one ``hitl-decision`` event in the same + // order as ``pendingInterrupts``. A ref because partial progress should not + // re-render the page. + const stagedDecisionsByInterruptIdRef = useRef>(new Map()); + const threadDetailQuery = useThreadDetail(activeThreadId); const threadMessagesQuery = useThreadMessages(activeThreadId); - // Live collaboration: sync session state and messages via Zero + // Live collaboration: sync session state and messages via Zero. Kept on the + // page because "AI responding" reflects the currently-viewed thread. useChatSessionStateSync(activeThreadId); const { data: membersData } = useAtomValue(membersAtom); @@ -498,10 +241,6 @@ export default function NewChatPage() { content: unknown; author_id: string | null; created_at: string; - // Forwarded so ``convertToThreadMessage`` can rebuild the - // ``metadata.custom.chatTurnId`` on the - // ``ThreadMessageLike``. Required by the inline Revert - // button's per-turn fallback. turn_id?: string | null; }[] ) => { @@ -520,7 +259,6 @@ export default function NewChatPage() { return reconcileInterruptedAssistantMessages(syncedMessages).map((msg) => { const member = msg.author_id ? (memberById.get(msg.author_id) ?? null) : null; - // Preserve existing author info if member lookup fails (e.g., cloned chats) const existingMsg = prevById.get(`msg-${msg.id}`); const existingAuthor = existingMsg?.metadata?.custom?.author as | { displayName?: string | null; avatarUrl?: string | null } @@ -535,10 +273,6 @@ export default function NewChatPage() { created_at: msg.created_at, author_display_name: member?.user_display_name ?? existingAuthor?.displayName ?? null, author_avatar_url: member?.user_avatar_url ?? existingAuthor?.avatarUrl ?? null, - // Forward the per-turn correlation id so the - // inline Revert button's ``(chat_turn_id, - // tool_name, position)`` fallback survives the - // post-stream Zero re-sync. turn_id: msg.turn_id ?? null, }); }); @@ -556,169 +290,33 @@ export default function NewChatPage() { return Number.isNaN(parsed) ? 0 : parsed; }, [params.workspace_id]); - // Unified store for agent-action rows (the same react-query cache - // the agent-actions dialog, the inline Revert button, and the - // per-turn Revert button all read). Hydrates from - // ``GET /threads/{id}/actions`` and is updated incrementally by the - // SSE handlers + revert-batch results below — no atom side-channel. + // Unified store for agent-action rows (react-query cache). Used by the + // edit pre-flight to count reversible downstream actions. const { items: agentActionItems } = useAgentActionsQuery(activeThreadId); - const handleChatFailure = useCallback( - async ({ - error, - flow, - threadId, - assistantMsgId, - }: { - error: unknown; - flow: ChatFlow; - threadId: number | null; - assistantMsgId: string; - }) => { - const normalized = classifyChatError({ - error, - flow, - context: { - workspaceId: workspaceId, - threadId, - }, - }); + // Latest displayed messages, read by the engine wrappers at call time so + // history/slice seeds stay fresh without re-creating the callbacks. + const messagesRef = useRef(displayMessages); + messagesRef.current = displayMessages; - const logger = - normalized.severity === "error" - ? console.error - : normalized.severity === "warn" - ? console.warn - : console.info; - logger(`[NewChatPage] ${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) { - setPremiumAlertForThread({ - 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); - }, - [currentUser?.id, persistAssistantErrorMessage, workspaceId, setPremiumAlertForThread] + const buildCtx = useCallback( + (): EngineContext => ({ + workspaceId, + threadId: activeThreadId, + priorMessages: messagesRef.current, + view: { setThreadId, setCurrentThread }, + }), + [workspaceId, activeThreadId] ); - const handleStreamTerminalError = useCallback( - async ({ - error, - flow, - threadId, - assistantMsgId, - accepted, - onAbort, - onPreAcceptFailure, - onAcceptedStreamError, - }: { - error: unknown; - flow: ChatFlow; - threadId: number | null; - assistantMsgId: string; - accepted: boolean; - onAbort?: () => Promise; - onPreAcceptFailure?: () => Promise; - onAcceptedStreamError?: () => 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", - }); - }, - [handleChatFailure] - ); - - const fetchWithTurnCancellingRetry = useCallback(async (runFetch: () => 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() - recentCancelRequestedAtRef.current <= 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", - }); - }, []); - const hydratedMessagesRef = useRef<{ threadId: number | null; data: typeof threadMessagesQuery.data; }>({ threadId: null, data: undefined }); - // Reset thread-local runtime state on route/workspace changes. Data fetching - // is handled by React Query below so the chat shell can render immediately. + // Reset thread-local runtime state on route/workspace changes. The durable + // streaming overlay is preserved for any still-running thread (and the newly + // viewed thread) via ``clearInactive`` so an in-flight turn survives nav. useEffect(() => { const nextThreadId = urlChatId > 0 ? urlChatId : null; handledLoadErrorThreadRef.current = null; @@ -732,13 +330,11 @@ export default function NewChatPage() { clearPlanOwnerRegistry(); closeReportPanel(); closeEditorPanel(); - // Note: agent-action data is keyed by threadId in react-query so - // switching threads naturally swaps caches; no explicit reset. + chatStreamStore.clearInactive(nextThreadId); }, [ urlChatId, setMentionedDocuments, setMessageDocumentsMap, - tokenUsageStore, closeReportPanel, closeEditorPanel, ]); @@ -773,6 +369,7 @@ export default function NewChatPage() { return; } + // Per-thread gate: never overwrite the live overlay of a running turn. if (isRunning) { return; } @@ -800,13 +397,18 @@ export default function NewChatPage() { } setMessageDocumentsMap(restoredDocsMap); hydratedMessagesRef.current = { threadId: activeThreadId, data: messagesResponse }; + + // The DB is now authoritative for this thread — drop the streaming + // overlay so we render DB messages (no-op while running / HITL-pending). + if (loadedMessages.length >= chatStreamStore.getMessages(activeThreadId).length) { + chatStreamStore.clear(activeThreadId); + } }, [ activeThreadId, isRunning, messages.length, setMessageDocumentsMap, threadMessagesQuery.data, - tokenUsageStore, ]); useEffect(() => { @@ -838,8 +440,7 @@ export default function NewChatPage() { threadMessagesQuery.error, ]); - // Prefetch document titles for @ mention picker - // Runs when user lands on page so data is ready when they type @ + // Prefetch document titles for @ mention picker so data is ready on type. useEffect(() => { if (!workspaceId) return; @@ -856,10 +457,7 @@ export default function NewChatPage() { }); }, [workspaceId, queryClient]); - // Handle scroll to comment from URL query params (e.g., from inbox item click) - // Read from window.location.search inside the effect instead of subscribing via - // useSearchParams() — avoids re-rendering this heavy component tree on every - // unrelated query-string change. (Vercel Best Practice: rerender-defer-reads 5.2) + // Handle scroll to comment from URL query params (e.g., from inbox click). useEffect(() => { const readAndApplyCommentId = () => { const params = new URLSearchParams(window.location.search); @@ -874,10 +472,8 @@ export default function NewChatPage() { readAndApplyCommentId(); - // Also respond to SPA navigations (back/forward) that change the query string window.addEventListener("popstate", readAndApplyCommentId); - // Cleanup on unmount or when navigating away return () => { window.removeEventListener("popstate", readAndApplyCommentId); clearTargetCommentId(); @@ -919,934 +515,158 @@ export default function NewChatPage() { setCurrentThreadMetadata, ]); - // Cleanup on unmount - abort any in-flight requests - useEffect(() => { - return () => { - if (abortControllerRef.current) { - abortControllerRef.current.abort(); - abortControllerRef.current = null; - } - }; - }, []); - - // Cancel ongoing request - const cancelRun = useCallback(async () => { - 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") { - recentCancelRequestedAtRef.current = Date.now(); - } - } - } catch (error) { - console.warn("[NewChatPage] Failed to signal cancel-active-turn:", error); - } - } - if (abortControllerRef.current) { - abortControllerRef.current.abort(); - abortControllerRef.current = null; - } - setIsRunning(false); - }, [threadId]); - // Handle new message from user const onNew = useCallback( + (message: AppendMessage) => startNewChat(buildCtx(), message), + [buildCtx] + ); + + // Cancel the in-flight turn (targets the active stream's owner thread). + const onCancel = useCallback(async () => { + await cancelActiveTurn(); + }, []); + + // Convert message (pass through since already in correct format) + const convertMessage = useCallback( + (message: ThreadMessageLike): ThreadMessageLike => message, + [] + ); + + // Handle editing a message - truncates history and regenerates with new + // query. When ``message.sourceId`` is set we pin ``from_message_id`` so the + // backend rewinds to the right checkpoint, and prompt the user to revert / + // continue / cancel before regenerating. + const onEdit = useCallback( async (message: AppendMessage) => { - // Abort any previous streaming request to prevent race conditions - // when user sends a second query while the first is still streaming - if (abortControllerRef.current) { - abortControllerRef.current.abort(); - abortControllerRef.current = null; + const { userQuery, userImages } = extractUserTurnForNewChatApi(message, []); + const queryForApi = userQuery.trim(); + if (!queryForApi && userImages.length === 0) { + toast.error("Cannot edit with empty message"); + return; } - // 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 activeMentions = submittedSnapshot ?? mentionedDocuments; - const mentionPayload = deriveMentionedPayload(activeMentions); - if (activeMentions.length > 0) { - setMentionedDocuments([]); + const userMessageContent = message.content as unknown as ThreadMessageLike["content"]; + + const sourceId = (message as { sourceId?: string }).sourceId; + const fromMessageId = + sourceId && /^msg-\d+$/.test(sourceId) + ? Number.parseInt(sourceId.replace(/^msg-/, ""), 10) + : null; + + if (fromMessageId == null) { + await regenerateChat(buildCtx(), queryForApi, { + userMessageContent, + userImages, + sourceUserMessageId: sourceId, + }); + return; } - const urlsSnapshot = [...pendingUserImageUrls]; - const { userQuery, userImages } = extractUserTurnForNewChatApi(message, urlsSnapshot); - - if (!userQuery.trim() && userImages.length === 0) 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; - setThreadId(currentThreadId); - // Set currentThread so share button in header appears immediately - setCurrentThread(newThread); - queryClient.setQueryData(cacheKeys.threads.detail(newThread.id), newThread); - queryClient.setQueryData(cacheKeys.threads.messages(newThread.id), { messages: [] }); - - // Track chat creation - 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("[NewChatPage] Failed to create thread:", error); - await handleChatFailure({ - error: tagPreAcceptSendFailure(error), - flow: "new", - threadId: currentThreadId, - assistantMsgId: "no-persist-assistant", - }); - return; + const msgs = messagesRef.current; + const editedIndex = msgs.findIndex((m) => m.id === `msg-${fromMessageId}`); + let downstreamReversibleCount = 0; + let downstreamTotalCount = 0; + if (editedIndex >= 0) { + const downstream = msgs.slice(editedIndex + 1); + downstreamTotalCount = downstream.length; + const seenTurns = new Set(); + const downstreamTurnIds = new Set(); + for (const m of downstream) { + const meta = (m.metadata ?? {}) as { custom?: { chatTurnId?: string } }; + const tid = meta.custom?.chatTurnId; + if (!tid || seenTurns.has(tid)) continue; + seenTurns.add(tid); + downstreamTurnIds.add(tid); + } + for (const a of agentActionItems) { + if (!a.chat_turn_id || !downstreamTurnIds.has(a.chat_turn_id)) continue; + if ( + a.reversible && + (a.reverted_by_action_id === null || a.reverted_by_action_id === undefined) && + !a.is_revert_action && + (a.error === null || a.error === undefined) + ) { + downstreamReversibleCount += 1; + } } } - if (urlsSnapshot.length > 0) { - setPendingUserImageUrls((prev) => prev.filter((u) => !urlsSnapshot.includes(u))); + if (downstreamReversibleCount === 0) { + await regenerateChat( + buildCtx(), + queryForApi, + { userMessageContent, userImages, sourceUserMessageId: sourceId }, + { fromMessageId, revertActions: false } + ); + return; } - // Add user message to state. Mutable because the SSE - // ``data-user-message-id`` handler (below) renames this - // optimistic id to the canonical ``msg-{db_id}`` once the - // backend's ``persist_user_turn`` resolves the row, and - // the in-stream flush / interrupt closures need to see - // the post-rename value via this live ``let`` binding. - let userMsgId = `msg-user-${Date.now()}`; - - // Always include author metadata so the UI layer can decide visibility - 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, - }; - setMessages((prev) => [...prev, userMessage]); - - // Track message sent - trackChatMessageSent(workspaceId, currentThreadId, { - hasAttachments: userImages.length > 0, - hasMentionedDocuments: - mentionPayload.document_ids.length > 0 || - mentionPayload.folder_ids.length > 0 || - mentionPayload.connector_ids.length > 0, - messageLength: userQuery.length, + setEditDialogState({ + fromMessageId, + userQuery: queryForApi, + userMessageContent, + userImages, + downstreamReversibleCount, + downstreamTotalCount, }); + }, + [buildCtx, agentActionItems] + ); - // Collect unique mention chips for display & persistence. - // The ``kind`` field is forwarded to the backend - // so the persisted ``mentioned-documents`` content part - // can render the correct chip type on reload. - 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); + const handleApprovalSubmit = useCallback( + (interruptId: string, decisions: HitlDecision[]) => { + // Stage this card's decisions; only fire the resume once every pending + // card in the current turn has submitted, so the backend slicer sees a + // single concatenated decisions list. + stagedDecisionsByInterruptIdRef.current.set(interruptId, decisions); + if (stagedDecisionsByInterruptIdRef.current.size < pendingInterrupts.length) { + return; } - - if (allMentionedDocs.length > 0) { - setMessageDocumentsMap((prev) => ({ - ...prev, - [userMsgId]: allMentionedDocs, - })); - } - - // Start streaming response - setIsRunning(true); - const controller = new AbortController(); - abortControllerRef.current = controller; - - // Prepare assistant message. Mutable for the same reason - // as ``userMsgId`` above — the ``data-assistant-message-id`` - // SSE handler reassigns this once - // ``persist_assistant_shell`` returns its canonical id. - 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 { - const selection = await getAgentFilesystemSelection(workspaceId, { - localFilesystemEnabled, - }); - 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."); + const ordered: HitlDecision[] = []; + for (const pi of pendingInterrupts) { + const staged = stagedDecisionsByInterruptIdRef.current.get(pi.interruptId); + if (!staged) { return; } - - // Build message history for context - const messageHistory = messages - .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); - - // Backend expects each mention kind in its own payload bucket. - 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: currentThreadId, - 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, - // Full mention metadata so the backend can persist a - // ``mentioned-documents`` ContentPart on the user message. - 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; - setMessages((prev) => [ - ...prev, - { - id: assistantMsgId, - role: "assistant", - content: [{ type: "text", text: "" }], - createdAt: new Date(), - }, - ]); - - const flushMessages = () => { - setMessages((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") { - recentCancelRequestedAtRef.current = Date.now(); - } - }, - }) - ) { - return; - } - switch (parsed.type) { - case "data-thread-title-update": { - const titleData = parsed.data as { threadId: number; title: string }; - if (titleData?.title && titleData?.threadId === currentThreadId) { - setCurrentThread((prev) => (prev ? { ...prev, title: titleData.title } : prev)); - updateChatTabTitle({ chatId: currentThreadId, 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) { - setAgentCreatedDocuments((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); - } - setMessages((prev) => - prev.map((m) => - m.id === assistantMsgId - ? { ...m, content: buildContentForUI(contentPartsState, toolsWithUI) } - : m - ) - ); - if (currentThreadId) { - // ``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: currentThreadId, - assistantMsgId, - interruptData, - bundleToolCallIds, - }; - setPendingInterrupts((prev) => { - const without = prev.filter((p) => p.interruptId !== interruptId); - return [...without, incoming]; - }); - } - } - break; - } - - case "data-action-log": { - applyActionLogSse(queryClient, currentThreadId, workspaceId, parsed.data); - break; - } - - case "data-action-log-updated": { - applyActionLogUpdatedSse( - queryClient, - currentThreadId, - parsed.data.id, - parsed.data.reversible - ); - break; - } - - case "data-turn-info": { - const turnId = readStreamedChatTurnId(parsed.data); - if (turnId) { - setMessages((prev) => - applyTurnIdToAssistantMessageList(prev, assistantMsgId, turnId) - ); - } - break; - } - - case "data-user-message-id": { - // Server-authoritative user message id resolved by - // ``persist_user_turn`` (or recovered via ON CONFLICT). - // Rename the optimistic ``msg-user-XXX`` placeholder to - // the canonical ``msg-{db_id}`` so DB-id-gated UI - // (comments, edit-from-this-message) unlocks immediately, - // migrate the local mentioned-documents map, and reassign - // the closure variable so all downstream - // ``m.id === userMsgId`` checks see the new value. - const parsedMsg = readStreamedMessageId(parsed.data); - if (!parsedMsg) break; - const newUserMsgId = `msg-${parsedMsg.messageId}`; - const oldUserMsgId = userMsgId; - setMessages((prev) => - prev.map((m) => - m.id === oldUserMsgId - ? mergeChatTurnIdIntoMessage({ ...m, id: newUserMsgId }, parsedMsg.turnId) - : m - ) - ); - if (allMentionedDocs.length > 0) { - setMessageDocumentsMap((prev) => { - if (!(oldUserMsgId in prev)) { - return { ...prev, [newUserMsgId]: allMentionedDocs }; - } - const { [oldUserMsgId]: _removed, ...rest } = prev; - return { ...rest, [newUserMsgId]: allMentionedDocs }; - }); - } - userMsgId = newUserMsgId; - if (isNewThread) { - // First user-side row landed in ``new_chat_messages``; - // refresh the sidebar so the freshly-bumped - // ``thread.updated_at`` reorders this thread. - queryClient.invalidateQueries({ - queryKey: ["threads", String(workspaceId)], - }); - } - break; - } - - case "data-assistant-message-id": { - // Server-authoritative assistant message id resolved - // by ``persist_assistant_shell``. Rename the optimistic - // id, migrate ``tokenUsageStore`` so any pending - // ``data-token-usage`` payload binds to the new id, - // remap any in-flight ``pendingInterrupts`` entries, - // and reassign the closure variable so the in-stream - // flush callback (line ~1074) keeps writing to the - // renamed message. - const parsedMsg = readStreamedMessageId(parsed.data); - if (!parsedMsg) break; - const newAssistantMsgId = `msg-${parsedMsg.messageId}`; - const oldAssistantMsgId = assistantMsgId; - tokenUsageStore.rename(oldAssistantMsgId, newAssistantMsgId); - setMessages((prev) => - prev.map((m) => - m.id === oldAssistantMsgId - ? mergeChatTurnIdIntoMessage({ ...m, id: newAssistantMsgId }, parsedMsg.turnId) - : m - ) - ); - setPendingInterrupts((prev) => - prev.map((p) => - p.assistantMsgId === oldAssistantMsgId - ? { ...p, assistantMsgId: newAssistantMsgId } - : p - ) - ); - assistantMsgId = newAssistantMsgId; - break; - } - } - }); - - batcher.flush(); - - // Server-authoritative persistence: ``stream_new_chat`` - // already wrote the user row in ``persist_user_turn`` - // (the FE renamed the optimistic id mid-stream via - // ``data-user-message-id``) and finalises the assistant - // row in ``finalize_assistant_turn`` from a shielded - // ``finally`` block. Nothing left for the FE to persist - // here — track the response and unblock the UI. - if (contentParts.length > 0 && !wasInterrupted) { - trackChatResponseReceived(workspaceId, currentThreadId); - } - } catch (error) { - streamBatcher?.dispose(); - await handleStreamTerminalError({ - error, - flow: "new", - threadId: currentThreadId, - assistantMsgId, - accepted: newAccepted, - // Server-side ``finalize_assistant_turn`` runs from a - // shielded ``anyio.CancelScope(shield=True)`` finally - // block, so partial content (incl. abort-mid-stream) - // is already persisted by the BE for the assistant - // row, and ``persist_user_turn`` ran before any LLM - // call. The FE's only remaining responsibility on - // abort / accepted-stream-error is to surface the - // error toast (handled by ``handleStreamTerminalError`` - // itself). - onPreAcceptFailure: async () => { - // Pre-accept failure means the BE never accepted the - // request — no server-side persistence ran. Roll - // back the optimistic UI insertions we made before - // the fetch so the user message and any local - // mentioned-docs metadata don't linger. - setMessages((prev) => prev.filter((m) => m.id !== userMsgId)); - setMessageDocumentsMap((prev) => { - if (!(userMsgId in prev)) return prev; - const { [userMsgId]: _removed, ...rest } = prev; - return rest; - }); - }, - }); - } finally { - setIsRunning(false); - abortControllerRef.current = null; - if (currentThreadId) { - void queryClient.invalidateQueries({ - queryKey: cacheKeys.threads.messages(currentThreadId), - }); - void queryClient.invalidateQueries({ - queryKey: cacheKeys.threads.detail(currentThreadId), - }); - } + ordered.push(...staged); } - }, - [ - threadId, - workspaceId, - messages, - jotaiStore, - mentionedDocuments, - setMentionedDocuments, - setMessageDocumentsMap, - setAgentCreatedDocuments, - queryClient, - currentUser, - localFilesystemEnabled, - disabledTools, - updateChatTabTitle, - tokenUsageStore, - pendingUserImageUrls, - setPendingUserImageUrls, - fetchWithTurnCancellingRetry, - handleStreamTerminalError, - handleChatFailure, - ] - ); - - const handleResume = useCallback( - async ( - decisions: Array<{ - type: string; - message?: string; - edited_action?: { name: string; args: Record }; - }> - ) => { - if (pendingInterrupts.length === 0) return; - // All cards in this turn share the same threadId/assistantMsgId - // (they're siblings of one parent agent step), so reading from - // the first entry is safe. - const resumeThreadId = pendingInterrupts[0].threadId; - // Destructured separately as ``let`` so the SSE - // ``data-assistant-message-id`` handler (resume always - // allocates a fresh server-side row) can rename it to - // the canonical ``msg-{db_id}`` mid-stream. - let assistantMsgId = pendingInterrupts[0].assistantMsgId; - // Concatenate every card's tool-call ids in pendingInterrupts order; - // this matches the ``decisions`` ordering produced by - // ``handleApprovalSubmit`` and the backend slicer's traversal of - // ``state.interrupts``. - const allBundleToolCallIds = pendingInterrupts.flatMap((p) => p.bundleToolCallIds); - setPendingInterrupts([]); stagedDecisionsByInterruptIdRef.current.clear(); - setIsRunning(true); - - const controller = new AbortController(); - abortControllerRef.current = controller; - - 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 = messages.find((m) => m.id === assistantMsgId); - if (existingMsg && Array.isArray(existingMsg.content)) { - // See ``ContentPartsState.suppressStepSeparators`` doc. - 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, - // argsText: assistant-ui prefers it over - // JSON.stringify(args), so restoring it keeps - // pretty-printed JSON across reloads. - ...(typeof p.argsText === "string" ? { argsText: p.argsText } : {}), - ...(typeof p.langchainToolCallId === "string" - ? { langchainToolCallId: p.langchainToolCallId } - : {}), - // metadata: spanId / thinkingStepId drive the - // timeline's step↔tool join. Dropping these - // here orphans every rehydrated tool-call. - ...(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) and multi-edit bundles do not - // collapse onto ``decisions[0]``. Cards outside the bundle are - // untouched. Mirrors the host ``hitl-decision`` handler. - 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; - // Sync argsText so the rendered card shows the - // edited inputs (assistant-ui prefers it over - // JSON.stringify(args)). - 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 = () => { - setMessages((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") { - recentCancelRequestedAtRef.current = Date.now(); - } - }, - }) - ) { - 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); - } - setMessages((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, - }; - setPendingInterrupts((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) { - setMessages((prev) => - applyTurnIdToAssistantMessageList(prev, assistantMsgId, turnId) - ); - } - break; - } - - case "data-assistant-message-id": { - // Resume always allocates a fresh ``new_chat_messages`` - // row anchored to a new ``turn_id`` (the original - // interrupted turn's row stays as-is), so this is a - // real id swap. Rename the optimistic placeholder to - // ``msg-{db_id}`` and reassign closure state. Resume - // does NOT emit ``data-user-message-id`` — the user - // row belongs to the original interrupted turn. - const parsedMsg = readStreamedMessageId(parsed.data); - if (!parsedMsg) break; - const newAssistantMsgId = `msg-${parsedMsg.messageId}`; - const oldAssistantMsgId = assistantMsgId; - tokenUsageStore.rename(oldAssistantMsgId, newAssistantMsgId); - setMessages((prev) => - prev.map((m) => - m.id === oldAssistantMsgId - ? mergeChatTurnIdIntoMessage({ ...m, id: newAssistantMsgId }, parsedMsg.turnId) - : m - ) - ); - assistantMsgId = newAssistantMsgId; - break; - } - } - }); - - batcher.flush(); - - // Server-authoritative persistence: ``stream_resume_chat`` - // finalises the assistant row in - // ``finalize_assistant_turn`` from a shielded - // ``finally`` block (covers both happy-path and - // abort-mid-stream). FE has no remaining persistence - // work here. - } catch (error) { - streamBatcher?.dispose(); - await handleStreamTerminalError({ - error, - flow: "resume", - threadId: resumeThreadId, - assistantMsgId, - accepted: resumeAccepted, - }); - } finally { - setIsRunning(false); - abortControllerRef.current = null; - void queryClient.invalidateQueries({ - queryKey: cacheKeys.threads.messages(resumeThreadId), - }); - void queryClient.invalidateQueries({ - queryKey: cacheKeys.threads.detail(resumeThreadId), - }); - } + window.dispatchEvent(new CustomEvent("hitl-decision", { detail: { decisions: ordered } })); }, - [ - pendingInterrupts, - messages, - workspaceId, - localFilesystemEnabled, - disabledTools, - queryClient, - tokenUsageStore, - fetchWithTurnCancellingRetry, - handleStreamTerminalError, - ] + [pendingInterrupts] ); + const handleEditDialogChoice = useCallback( + async (choice: EditMessageDialogChoice) => { + const pending = editDialogState; + if (!pending) return; + setEditDialogState(null); + if (choice === "cancel") return; + await regenerateChat( + buildCtx(), + pending.userQuery, + { + userMessageContent: pending.userMessageContent, + userImages: pending.userImages, + sourceUserMessageId: `msg-${pending.fromMessageId}`, + }, + { + fromMessageId: pending.fromMessageId, + revertActions: choice === "revert", + } + ); + }, + [editDialogState, buildCtx] + ); + + // Handle reloading/refreshing the last AI response + const onReload = useCallback(async () => { + await regenerateChat(buildCtx(), null); + }, [buildCtx]); + + // HITL resume bridge. Submit always happens from this page's approval UI, so + // the currently-viewed thread owns the pending interrupts. Applies each + // decision to its card, then resumes the (durable) stream. useEffect(() => { const handler = (e: Event) => { const detail = (e as CustomEvent).detail as { @@ -1859,15 +679,9 @@ export default function NewChatPage() { if (!detail?.decisions || pendingInterrupts.length === 0) return; const incoming = detail.decisions; if (incoming.length === 0) return; - // Concatenated tool-call ids across every pending card, in the - // order ``handleApprovalSubmit`` produced ``incoming``. const tcIds = pendingInterrupts.flatMap((p) => p.bundleToolCallIds); const N = tcIds.length; - // Refuse rather than silently broadcast or drop. The orchestrator - // only fires ``hitl-decision`` once every pending card has - // submitted, so a count mismatch indicates a contract drift - // (and would later make the backend slicer raise). if (incoming.length !== N) { toast.error( `Cannot resume: ${incoming.length} decision(s) submitted for ${N} pending actions.` @@ -1890,631 +704,71 @@ export default function NewChatPage() { submittedDecisions.push(decision); } - // All pending cards belong to the same assistant message, so a - // single content-update pass suffices. const targetAssistantMsgId = pendingInterrupts[0].assistantMsgId; - setMessages((prev) => - prev.map((m) => { - if (m.id !== targetAssistantMsgId) return m; - const parts = m.content as unknown as Array>; - const newContent = parts.map((part) => { - const tcId = part.toolCallId as string | undefined; - const d = tcId ? byTcId.get(tcId) : undefined; - if (!d || part.type !== "tool-call") return part; - if (typeof part.result !== "object" || part.result === null) return part; - if (!("__interrupt__" in (part.result as Record))) return part; - const decided = d.type; - if (decided === "edit" && d.edited_action) { + if (activeThreadId != null) { + chatStreamStore.setMessages(activeThreadId, (prev) => + prev.map((m) => { + if (m.id !== targetAssistantMsgId) return m; + const parts = m.content as unknown as Array>; + const newContent = parts.map((part) => { + const tcId = part.toolCallId as string | undefined; + const d = tcId ? byTcId.get(tcId) : undefined; + if (!d || part.type !== "tool-call") return part; + if (typeof part.result !== "object" || part.result === null) return part; + if (!("__interrupt__" in (part.result as Record))) return part; + const decided = d.type; + if (decided === "edit" && d.edited_action) { + return { + ...part, + args: d.edited_action.args, + argsText: JSON.stringify(d.edited_action.args, null, 2), + result: { + ...(part.result as Record), + __decided__: decided, + }, + }; + } return { ...part, - args: d.edited_action.args, - // Sync argsText so the card renders the edited - // inputs (assistant-ui prefers it over JSON.stringify). - argsText: JSON.stringify(d.edited_action.args, null, 2), result: { ...(part.result as Record), __decided__: decided, }, }; - } - return { - ...part, - result: { - ...(part.result as Record), - __decided__: decided, - }, - }; - }); - return { ...m, content: newContent as unknown as ThreadMessageLike["content"] }; - }) - ); - handleResume(submittedDecisions); + }); + return { ...m, content: newContent as unknown as ThreadMessageLike["content"] }; + }) + ); + } + void resumeChat(buildCtx(), submittedDecisions); }; window.addEventListener("hitl-decision", handler); return () => window.removeEventListener("hitl-decision", handler); - }, [handleResume, pendingInterrupts]); - - // Convert message (pass through since already in correct format) - const convertMessage = useCallback( - (message: ThreadMessageLike): ThreadMessageLike => message, - [] - ); - - /** - * Handle regeneration (edit or reload) by calling the regenerate endpoint - * and streaming the response. This rewinds the LangGraph checkpointer state. - * - * @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, - editExtras?: { - userMessageContent: ThreadMessageLike["content"]; - userImages: NewChatUserImagePayload[]; - sourceUserMessageId?: string; - }, - editFromPosition?: { - /** Message id (numeric, parsed from ``msg-``) to rewind to. */ - fromMessageId?: number | null; - /** When true, revert reversible downstream actions before stream. */ - revertActions?: boolean; - } - ) => { - 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(); - abortControllerRef.current = null; - } - - // Extract the original user query BEFORE removing messages (for 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) { - // Reload mode - find and preserve the last user message content - const lastUserMessage = [...messages].reverse().find((m) => m.role === "user"); - if (lastUserMessage) { - sourceUserMessageId = lastUserMessage.id; - originalUserMessageContent = lastUserMessage.content; - originalUserMessageMetadata = lastUserMessage.metadata; - // Extract text for the API request - for (const part of lastUserMessage.content) { - if (typeof part === "object" && part.type === "text" && "text" in part) { - userQueryToDisplay = part.text; - break; - } - } - } - } else { - userQueryToDisplay = newUserQuery; - } - - // Start streaming - setIsRunning(true); - const controller = new AbortController(); - abortControllerRef.current = controller; - - // Add placeholder user message if we have a new query (edit mode). - // Mutable for the same reason as in ``onNew`` — both ids are - // renamed mid-stream by the new ``data-user-message-id`` / - // ``data-assistant-message-id`` SSE handlers below. - 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; - - // Add placeholder messages to UI - // Always add back the user message (with new query for edit, or original content for reload) - 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, - }); - // Partition the source mentions back into doc/folder id buckets - // so the regenerate route can pass them to ``stream_new_chat`` - // and the priority middleware sees the same ``[USER-MENTIONED]`` - // priority entries the original turn did. Without this partition - // the regenerate flow silently dropped the agent's mention - // awareness — same architectural bug we fixed on the new-chat path. - 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, - // Full mention metadata for the regenerate-specific - // source list. Only meaningful for edit (the BE only - // re-persists a user row when ``user_query`` is set); - // reload reuses the original turn's mentioned_documents. - mentioned_documents: sourceMentionedDocs.length > 0 ? sourceMentionedDocs : undefined, - }; - if (isEdit) { - requestBody.user_images = editExtras?.userImages ?? []; - } - // Explicit edit-from-arbitrary-position. Only send - // ``from_message_id`` / ``revert_actions`` when the - // caller asked for them; otherwise the backend keeps the - // legacy "last 2 messages" behaviour for back-compat. - if (editFromPosition?.fromMessageId != null) { - requestBody.from_message_id = editFromPosition.fromMessageId; - if (editFromPosition.revertActions) { - requestBody.revert_actions = true; - } - } - const response = await fetchWithTurnCancellingRetry(() => - authenticatedFetch(getRegenerateUrl(threadId), { - method: "POST", - headers: { - "Content-Type": "application/json", - }, - body: JSON.stringify(requestBody), - signal: controller.signal, - }) - ); - - if (!response.ok) { - throw await toHttpResponseError(response); - } - regenerateAccepted = true; - - // Only switch UI to regenerated placeholder messages after the backend accepts - // regenerate. This avoids local message loss when regenerate fails early (e.g. 400). - // - // When an explicit ``editFromPosition.fromMessageId`` is passed, slice from - // that message forward so edit-from-arbitrary-position drops every downstream - // message; otherwise fall back to the legacy "drop the last 2" behaviour. - setMessages((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) { - setMessageDocumentsMap((prev) => ({ - ...prev, - [userMsgId]: sourceMentionedDocs, - })); - } - - const flushMessages = () => { - setMessages((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") { - recentCancelRequestedAtRef.current = Date.now(); - } - }, - }) - ) { - return; - } - switch (parsed.type) { - case "data-action-log": { - if (threadId !== null) { - applyActionLogSse(queryClient, threadId, workspaceId, parsed.data); - } - break; - } - - case "data-action-log-updated": { - if (threadId !== null) { - applyActionLogUpdatedSse( - queryClient, - threadId, - parsed.data.id, - parsed.data.reversible - ); - } - break; - } - - case "data-turn-info": { - const turnId = readStreamedChatTurnId(parsed.data); - if (turnId) { - setMessages((prev) => - applyTurnIdToAssistantMessageList(prev, assistantMsgId, turnId) - ); - } - break; - } - - case "data-user-message-id": { - // Same role as in ``onNew`` but the regenerate-specific - // mention metadata (``sourceMentionedDocs``) is the - // list to migrate onto the canonical id key. - const parsedMsg = readStreamedMessageId(parsed.data); - if (!parsedMsg) break; - const newUserMsgId = `msg-${parsedMsg.messageId}`; - const oldUserMsgId = userMsgId; - setMessages((prev) => - prev.map((m) => - m.id === oldUserMsgId - ? mergeChatTurnIdIntoMessage({ ...m, id: newUserMsgId }, parsedMsg.turnId) - : m - ) - ); - if (sourceMentionedDocs.length > 0) { - setMessageDocumentsMap((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); - setMessages((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; - // failureCount must include every "not undone" bucket - // (not_reversible, permission_denied, failed) so the - // toast's "X could not be rolled back" math matches - // the response invariant ``total === sum(counters)``. - // ``skipped`` rows are batch revert artefacts (revert - // rows themselves) and are not user-facing failures. - 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.` - ); - } - if (threadId !== null) { - for (const r of summary.results) { - if (r.status === "reverted" || r.status === "already_reverted") { - markActionRevertedInCache( - queryClient, - threadId, - r.action_id, - r.new_action_id ?? null - ); - } - } - } - break; - } - } - }); - - batcher.flush(); - - // Server-authoritative persistence: ``stream_new_chat`` - // (regenerate flow) wrote the user row in - // ``persist_user_turn`` and finalises the assistant row - // in ``finalize_assistant_turn`` from a shielded - // ``finally`` block (covers both happy-path and - // abort-mid-stream). FE only needs to track the - // successful response here. - if (contentParts.length > 0) { - trackChatResponseReceived(workspaceId, threadId); - } - } catch (error) { - streamBatcher?.dispose(); - await handleStreamTerminalError({ - error, - flow: "regenerate", - threadId, - assistantMsgId, - accepted: regenerateAccepted, - }); - } finally { - setIsRunning(false); - abortControllerRef.current = null; - void queryClient.invalidateQueries({ - queryKey: cacheKeys.threads.messages(threadId), - }); - void queryClient.invalidateQueries({ - queryKey: cacheKeys.threads.detail(threadId), - }); - } - }, - [ - threadId, - workspaceId, - messages, - disabledTools, - localFilesystemEnabled, - messageDocumentsMap, - setMessageDocumentsMap, - queryClient, - tokenUsageStore, - fetchWithTurnCancellingRetry, - handleStreamTerminalError, - ] - ); - - // Handle editing a message - truncates history and regenerates with new query. - // - // When ``message.sourceId`` is set (the assistant-ui way to say - // "this edit replaces an older message"), we pin - // ``from_message_id`` so the backend rewinds to the right LangGraph - // checkpoint instead of relying on the legacy "last 2 messages" - // rewind. We also count downstream reversible actions and prompt the - // user to revert / continue / cancel before regenerating. - const onEdit = useCallback( - async (message: AppendMessage) => { - const { userQuery, userImages } = extractUserTurnForNewChatApi(message, []); - const queryForApi = userQuery.trim(); - if (!queryForApi && userImages.length === 0) { - toast.error("Cannot edit with empty message"); - return; - } - - const userMessageContent = message.content as unknown as ThreadMessageLike["content"]; - - // ``sourceId`` per @assistant-ui/core's ``AppendMessage`` is - // "the ID of the message that was edited". Parse the numeric - // suffix so we can map it back to a DB row. - const sourceId = (message as { sourceId?: string }).sourceId; - const fromMessageId = - sourceId && /^msg-\d+$/.test(sourceId) - ? Number.parseInt(sourceId.replace(/^msg-/, ""), 10) - : null; - - if (fromMessageId == null) { - // No source id (or non-DB id) — fall back to today's - // last-2 behaviour. The user gets the legacy edit flow. - await handleRegenerate(queryForApi, { - userMessageContent, - userImages, - sourceUserMessageId: sourceId, - }); - return; - } - - // Pre-flight: count reversible downstream actions so we can - // auto-skip the dialog for harmless edits. - // - // "Downstream" means messages AFTER the edited one. The - // previous slice ``messages.slice(editedIndex)`` included - // the edited message itself in both the total - // count and the reversibility scan (any actions on the - // edited turn would be double-counted). Slice from - // ``editedIndex + 1`` so the dialog text matches reality: - // "N downstream messages will be dropped". - const editedIndex = messages.findIndex((m) => m.id === `msg-${fromMessageId}`); - let downstreamReversibleCount = 0; - let downstreamTotalCount = 0; - if (editedIndex >= 0) { - const downstream = messages.slice(editedIndex + 1); - downstreamTotalCount = downstream.length; - const seenTurns = new Set(); - const downstreamTurnIds = new Set(); - for (const m of downstream) { - const meta = (m.metadata ?? {}) as { custom?: { chatTurnId?: string } }; - const tid = meta.custom?.chatTurnId; - if (!tid || seenTurns.has(tid)) continue; - seenTurns.add(tid); - downstreamTurnIds.add(tid); - } - // Source of truth: the unified react-query cache. Every - // action whose ``chat_turn_id`` belongs to the slice we're - // about to drop counts toward the prompt. - for (const a of agentActionItems) { - if (!a.chat_turn_id || !downstreamTurnIds.has(a.chat_turn_id)) continue; - if ( - a.reversible && - (a.reverted_by_action_id === null || a.reverted_by_action_id === undefined) && - !a.is_revert_action && - (a.error === null || a.error === undefined) - ) { - downstreamReversibleCount += 1; - } - } - } - - if (downstreamReversibleCount === 0) { - // Nothing to revert — submit silently. - await handleRegenerate( - queryForApi, - { userMessageContent, userImages, sourceUserMessageId: sourceId }, - { fromMessageId, revertActions: false } - ); - return; - } - - setEditDialogState({ - fromMessageId, - userQuery: queryForApi, - userMessageContent, - userImages, - downstreamReversibleCount, - downstreamTotalCount, - }); - }, - [handleRegenerate, messages, agentActionItems] - ); - - const handleApprovalSubmit = useCallback( - (interruptId: string, decisions: HitlDecision[]) => { - // Stage this card's decisions; only fire the resume once every - // pending card in the current turn has submitted, so the - // backend slicer sees a single concatenated decisions list - // whose total matches the parent state's pending action count. - stagedDecisionsByInterruptIdRef.current.set(interruptId, decisions); - if (stagedDecisionsByInterruptIdRef.current.size < pendingInterrupts.length) { - return; - } - const ordered: HitlDecision[] = []; - for (const pi of pendingInterrupts) { - const staged = stagedDecisionsByInterruptIdRef.current.get(pi.interruptId); - if (!staged) { - // Defensive: a missing entry means the staging map and - // the pending list disagreed for one cycle. Bail rather - // than dispatch a count-mismatched batch. - return; - } - ordered.push(...staged); - } - stagedDecisionsByInterruptIdRef.current.clear(); - window.dispatchEvent(new CustomEvent("hitl-decision", { detail: { decisions: ordered } })); - }, - [pendingInterrupts] - ); - - const handleEditDialogChoice = useCallback( - async (choice: EditMessageDialogChoice) => { - const pending = editDialogState; - if (!pending) return; - setEditDialogState(null); - if (choice === "cancel") return; - await handleRegenerate( - pending.userQuery, - { - userMessageContent: pending.userMessageContent, - userImages: pending.userImages, - sourceUserMessageId: `msg-${pending.fromMessageId}`, - }, - { - fromMessageId: pending.fromMessageId, - revertActions: choice === "revert", - } - ); - }, - [editDialogState, handleRegenerate] - ); - - // Handle reloading/refreshing the last AI response - const onReload = useCallback(async () => { - // parentId is the ID of the message to reload from (the user message) - // We call regenerate without a query to use the same query - await handleRegenerate(null); - }, [handleRegenerate]); + }, [buildCtx, pendingInterrupts, activeThreadId]); // Surface the thread's deliverables to the layout-level artifacts sidebar. - useSyncChatArtifacts(messages); + useSyncChatArtifacts(displayMessages); // Create external store runtime const runtime = useExternalStoreRuntime({ - messages, + messages: displayMessages, isRunning, onNew, onEdit, onReload, convertMessage, - onCancel: cancelRun, + onCancel, }); const threadLoadError = activeThreadId ? (threadDetailQuery.error ?? threadMessagesQuery.error) : null; const shouldShowThreadLoadError = - !!threadLoadError && !!activeThreadId && !currentThread && messages.length === 0; + !!threadLoadError && !!activeThreadId && !currentThread && displayMessages.length === 0; const isThreadMessagesLoading = !!activeThreadId && threadMessagesQuery.isPending && - messages.length === 0 && + displayMessages.length === 0 && !threadMessagesQuery.error; if (shouldShowThreadLoadError) { diff --git a/surfsense_web/app/dashboard/[workspace_id]/user-settings/appearance/page.tsx b/surfsense_web/app/dashboard/[workspace_id]/user-settings/appearance/page.tsx new file mode 100644 index 000000000..6ae0952d5 --- /dev/null +++ b/surfsense_web/app/dashboard/[workspace_id]/user-settings/appearance/page.tsx @@ -0,0 +1,5 @@ +import { AppearanceContent } from "../components/AppearanceContent"; + +export default function Page() { + return ; +} diff --git a/surfsense_web/app/dashboard/[workspace_id]/user-settings/components/AppearanceContent.tsx b/surfsense_web/app/dashboard/[workspace_id]/user-settings/components/AppearanceContent.tsx new file mode 100644 index 000000000..258a31dfe --- /dev/null +++ b/surfsense_web/app/dashboard/[workspace_id]/user-settings/components/AppearanceContent.tsx @@ -0,0 +1,43 @@ +"use client"; + +import { useAtom } from "jotai"; +import { showMessageTimestampsAtom } from "@/atoms/chat/show-timestamps.atom"; +import { Label } from "@/components/ui/label"; +import { Switch } from "@/components/ui/switch"; + +export function AppearanceContent() { + const [showTimestamps, setShowTimestamps] = useAtom(showMessageTimestampsAtom); + + return ( +
+
+
+

Chat

+

+ Control how messages are displayed in your conversations. +

+
+
+
+
+ +

+ Display the time under each message in a chat. Saved on this device. +

+
+ +
+
+
+
+ ); +} diff --git a/surfsense_web/app/dashboard/[workspace_id]/user-settings/layout-shell.tsx b/surfsense_web/app/dashboard/[workspace_id]/user-settings/layout-shell.tsx index aa73917ef..b18db9e64 100644 --- a/surfsense_web/app/dashboard/[workspace_id]/user-settings/layout-shell.tsx +++ b/surfsense_web/app/dashboard/[workspace_id]/user-settings/layout-shell.tsx @@ -7,6 +7,7 @@ import { Library, MessageCircle, Monitor, + Palette, ReceiptText, ShieldCheck, WandSparkles, @@ -20,6 +21,7 @@ import { usePlatform } from "@/hooks/use-platform"; export type UserSettingsTab = | "profile" + | "appearance" | "api-key" | "prompts" | "community-prompts" @@ -49,6 +51,12 @@ export function UserSettingsLayoutShell({ workspaceId, children }: UserSettingsL href: `/dashboard/${workspaceId}/user-settings/profile`, icon: , }, + { + value: "appearance" as const, + label: "Appearance", + href: `/dashboard/${workspaceId}/user-settings/appearance`, + icon: , + }, { value: "api-key" as const, label: t("api_key_nav_label"), diff --git a/surfsense_web/atoms/chat/show-timestamps.atom.ts b/surfsense_web/atoms/chat/show-timestamps.atom.ts new file mode 100644 index 000000000..434f065ff --- /dev/null +++ b/surfsense_web/atoms/chat/show-timestamps.atom.ts @@ -0,0 +1,11 @@ +import { atomWithStorage } from "jotai/utils"; + +/** + * Per-device preference: show a timestamp under each chat message. + * + * Off by default to match streaming-AI chat convention (ChatGPT/Claude keep + * the message stream clean and put time in the conversation list). Persisted + * in localStorage, so it does not sync across devices — acceptable for a + * cosmetic display toggle. + */ +export const showMessageTimestampsAtom = atomWithStorage("chat-show-timestamps:v1", false); diff --git a/surfsense_web/components/assistant-ui/assistant-message.tsx b/surfsense_web/components/assistant-ui/assistant-message.tsx index c4a0e86dc..a45c651b1 100644 --- a/surfsense_web/components/assistant-ui/assistant-message.tsx +++ b/surfsense_web/components/assistant-ui/assistant-message.tsx @@ -35,6 +35,7 @@ import { useAllCitationMetadata, } from "@/components/assistant-ui/citation-metadata-context"; import { MarkdownText } from "@/components/assistant-ui/markdown-text"; +import { MessageTimestamp } from "@/components/assistant-ui/message-timestamp"; import { ReasoningMessagePart } from "@/components/assistant-ui/reasoning-message-part"; import { RevertTurnButton } from "@/components/assistant-ui/revert-turn-button"; import { @@ -62,6 +63,7 @@ import { withArtifactAnchor } from "@/features/chat-artifacts"; import { useComments } from "@/hooks/use-comments"; import { useMediaQuery } from "@/hooks/use-media-query"; import { useElectronAPI } from "@/hooks/use-platform"; +import { formatMessageTimestamp } from "@/lib/format-date"; import { getProviderIcon } from "@/lib/provider-icons"; import { tryGetHostname } from "@/lib/url"; import { cn } from "@/lib/utils"; @@ -249,16 +251,6 @@ export const MessageError: FC = () => { ); }; -function formatMessageDate(date: Date): string { - return date.toLocaleDateString(undefined, { - month: "short", - day: "numeric", - hour: "numeric", - minute: "2-digit", - hour12: true, - }); -} - /** * Format provider USD cost (in micro-USD) for inline display next to a * token count. Falls back to ``"<$0.001"`` for sub-tenth-of-a-cent @@ -367,7 +359,7 @@ const MessageInfoDropdown: FC<{ chatTurnId: string | null | undefined }> = ({ ch > {createdAt && ( - {formatMessageDate(createdAt)} + {formatMessageTimestamp(createdAt)} )} {hasUsage && ( @@ -463,6 +455,8 @@ const AssistantMessageInner: FC = () => { + + {isMobile && (
diff --git a/surfsense_web/components/assistant-ui/message-timestamp.tsx b/surfsense_web/components/assistant-ui/message-timestamp.tsx new file mode 100644 index 000000000..65a1da9be --- /dev/null +++ b/surfsense_web/components/assistant-ui/message-timestamp.tsx @@ -0,0 +1,24 @@ +import { useAuiState } from "@assistant-ui/react"; +import { useAtomValue } from "jotai"; +import type { FC } from "react"; +import { showMessageTimestampsAtom } from "@/atoms/chat/show-timestamps.atom"; +import { formatMessageTimestamp } from "@/lib/format-date"; +import { cn } from "@/lib/utils"; + +/** + * Muted, always-visible timestamp under a chat message. Renders only when the + * user has opted in via {@link showMessageTimestampsAtom} and the message + * carries a ``createdAt`` (absent on optimistic pre-persist messages). + */ +export const MessageTimestamp: FC<{ className?: string }> = ({ className }) => { + const show = useAtomValue(showMessageTimestampsAtom); + const createdAt = useAuiState(({ message }) => message?.createdAt); + + if (!show || !createdAt) return null; + + return ( +
+ {formatMessageTimestamp(createdAt)} +
+ ); +}; diff --git a/surfsense_web/components/assistant-ui/user-message.tsx b/surfsense_web/components/assistant-ui/user-message.tsx index 592dc4224..09a700ff1 100644 --- a/surfsense_web/components/assistant-ui/user-message.tsx +++ b/surfsense_web/components/assistant-ui/user-message.tsx @@ -22,6 +22,7 @@ import { currentThreadAtom } from "@/atoms/chat/current-thread.atom"; import { messageDocumentsMapAtom } from "@/atoms/chat/mentioned-documents.atom"; import { openEditorPanelAtom } from "@/atoms/editor/editor-panel.atom"; import { MentionChip } from "@/components/assistant-ui/mention-chip"; +import { MessageTimestamp } from "@/components/assistant-ui/message-timestamp"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; import { getConnectorIcon } from "@/contracts/enums/connectorIcons"; import { getMentionDocKey } from "@/lib/chat/mention-doc-key"; @@ -182,6 +183,7 @@ export const UserMessage: FC = () => {
)} + ); diff --git a/surfsense_web/components/chat/active-chat-stream-runner.tsx b/surfsense_web/components/chat/active-chat-stream-runner.tsx new file mode 100644 index 000000000..dd27922a0 --- /dev/null +++ b/surfsense_web/components/chat/active-chat-stream-runner.tsx @@ -0,0 +1,23 @@ +"use client"; + +import { useEffect } from "react"; +import { chatStreamStore } from "@/lib/chat/stream-engine/store"; + +/** + * Persistent, render-null host that scopes the in-flight chat turn's lifetime + * to the workspace shell, not the chat page. + * + * Mounted in ``LayoutDataProvider``, it survives in-app navigation between + * workspace routes and aborts the single active turn only on workspace/app + * teardown, so ordinary navigation disconnects the view without stopping the + * stream. + */ +export function ActiveChatStreamRunner() { + useEffect(() => { + return () => { + chatStreamStore.abortActive(); + }; + }, []); + + return null; +} diff --git a/surfsense_web/components/layout/providers/LayoutDataProvider.tsx b/surfsense_web/components/layout/providers/LayoutDataProvider.tsx index f9c2c9072..81f0d974f 100644 --- a/surfsense_web/components/layout/providers/LayoutDataProvider.tsx +++ b/surfsense_web/components/layout/providers/LayoutDataProvider.tsx @@ -18,6 +18,7 @@ import { workspacesAtom } from "@/atoms/workspaces/workspace-query.atoms"; import { ActionLogDialog } from "@/components/agent-action-log/action-log-dialog"; import { AnnouncementSpotlight } from "@/components/announcements/AnnouncementSpotlight"; import { AnnouncementsDialog } from "@/components/announcements/AnnouncementsDialog"; +import { ActiveChatStreamRunner } from "@/components/chat/active-chat-stream-runner"; import { AlertDialog, AlertDialogAction, @@ -644,6 +645,9 @@ export function LayoutDataProvider({ workspaceId, children }: LayoutDataProvider return ( <> + {/* Persistent host: keeps an in-flight chat turn streaming across + in-app navigation and aborts it only on workspace teardown. */} + >; + 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 +// --------------------------------------------------------------------------- + +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 }; diff --git a/surfsense_web/lib/chat/stream-engine/helpers.ts b/surfsense_web/lib/chat/stream-engine/helpers.ts new file mode 100644 index 000000000..afa1851e1 --- /dev/null +++ b/surfsense_web/lib/chat/stream-engine/helpers.ts @@ -0,0 +1,154 @@ +import { z } from "zod"; +import type { MentionedDocumentInfo } from "@/atoms/chat/mentioned-documents.atom"; +import type { ToolUIGate } from "@/lib/chat/streaming-state"; + +/** + * Every tool call renders a card. The sentinel ``"all"`` matches every tool + * — the legacy ``BASE_TOOLS_WITH_UI`` allowlist was dropped so unknown tool + * calls route through the generic ``ToolFallback``. Persisted payload size + * stays bounded because the backend's ``format_thinking_step`` summarisation + * and the ``result_length``-only default for unknown tools keep the JSON + * from ballooning. + */ +export const TOOLS_WITH_UI_ALL: ToolUIGate = "all"; + +export const TURN_CANCELLING_INITIAL_DELAY_MS = 200; +export const TURN_CANCELLING_BACKOFF_FACTOR = 2; +export const TURN_CANCELLING_MAX_DELAY_MS = 1500; +export const RECENT_CANCEL_WINDOW_MS = 5_000; + +export function sleep(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +export function computeFallbackTurnCancellingRetryDelay(attempt: number): number { + const safeAttempt = Math.max(1, attempt); + const raw = + TURN_CANCELLING_INITIAL_DELAY_MS * TURN_CANCELLING_BACKOFF_FACTOR ** (safeAttempt - 1); + return Math.min(raw, TURN_CANCELLING_MAX_DELAY_MS); +} + +/** + * Generate a synthetic ``toolCallId`` for an action_request that has no + * matching streamed tool-call card (HITL-blocked subagent calls don't surface + * as tool-call events). Suffixes a counter when the base id is already taken + * — sequential interrupts for the same tool name otherwise collide on + * ``interrupt-${name}-${i}`` and crash assistant-ui with a duplicate-key error. + */ +export function freshSynthToolCallId( + toolCallIndices: Map, + toolName: string, + index: number +): string { + const base = `interrupt-${toolName}-${index}`; + if (!toolCallIndices.has(base)) return base; + let n = 1; + while (toolCallIndices.has(`${base}-${n}`)) n++; + return `${base}-${n}`; +} + +/** + * Pair each ``action_request`` to a unique pending tool-call card, preserving + * order so ``decisions[i]`` lines up with ``action_requests[i]`` on the wire. + * + * Same-name bundles (e.g. three ``create_jira_issue``) used to collapse onto + * one card because the matcher keyed by name; this consumes each card via the + * ``claimed`` set and walks forward in DOM order. + */ +export function pairBundleToolCallIds( + toolCallIndices: Map, + contentParts: Array<{ + type: string; + toolName?: string; + result?: unknown; + }>, + actionRequests: ReadonlyArray<{ name: string }> +): Array { + const claimed = new Set(); + const paired: Array = []; + for (const action of actionRequests) { + let matched: string | null = null; + for (const [tcId, idx] of toolCallIndices) { + if (claimed.has(tcId)) continue; + const part = contentParts[idx]; + if (!part || part.type !== "tool-call" || part.toolName !== action.name) continue; + const result = part.result as Record | undefined | null; + if (result == null || (result.__interrupt__ === true && !result.__decided__)) { + matched = tcId; + claimed.add(tcId); + break; + } + } + paired.push(matched); + } + return paired; +} + +/** + * Zod schema for mentioned document info (for type-safe parsing). + * + * ``kind`` defaults to ``"doc"`` so messages persisted before folder + * mentions existed deserialise unchanged. + */ +const MentionedDocumentInfoSchema = z.object({ + id: z.number(), + title: z.string(), + document_type: z.string().optional(), + kind: z + .union([z.literal("doc"), z.literal("folder"), z.literal("connector"), z.literal("thread")]) + .optional() + .default("doc"), + connector_type: z.string().optional(), + account_name: z.string().optional(), +}); + +const MentionedDocumentsPartSchema = z.object({ + type: z.literal("mentioned-documents"), + documents: z.array(MentionedDocumentInfoSchema), +}); + +/** + * Extract mentioned documents from message content (type-safe with Zod). + */ +export function extractMentionedDocuments(content: unknown): MentionedDocumentInfo[] { + if (!Array.isArray(content)) return []; + + for (const part of content) { + const result = MentionedDocumentsPartSchema.safeParse(part); + if (result.success) { + return result.data.documents.map((doc) => { + if (doc.kind === "connector") { + return { + id: doc.id, + title: doc.title, + kind: "connector", + connector_type: doc.connector_type ?? doc.document_type ?? "UNKNOWN", + account_name: doc.account_name ?? doc.title, + }; + } + if (doc.kind === "folder") { + return { + id: doc.id, + title: doc.title, + kind: "folder", + }; + } + if (doc.kind === "thread") { + return { + id: doc.id, + title: doc.title, + kind: "thread", + }; + } + return { + id: doc.id, + title: doc.title, + document_type: doc.document_type ?? "UNKNOWN", + kind: "doc", + }; + }); + } + } + + return []; +} diff --git a/surfsense_web/lib/chat/stream-engine/store.ts b/surfsense_web/lib/chat/stream-engine/store.ts new file mode 100644 index 000000000..44627663f --- /dev/null +++ b/surfsense_web/lib/chat/stream-engine/store.ts @@ -0,0 +1,177 @@ +import type { ThreadMessageLike } from "@assistant-ui/react"; +import { createTokenUsageStore } from "@/components/assistant-ui/token-usage-context"; +import type { PendingInterruptState } from "@/features/chat-messages/hitl"; + +/** + * Durable, per-thread streaming state for a single in-flight chat turn. + * + * Lives at module scope (see {@link chatStreamStore}) so a running turn's + * messages / interrupts survive the chat page unmounting during in-app + * navigation. The React tree consumes it through ``useChatStream`` via + * ``useSyncExternalStore`` (state lives outside the render cycle). + */ +export interface ThreadStreamState { + threadId: number; + messages: ThreadMessageLike[]; + isRunning: boolean; + pendingInterrupts: PendingInterruptState[]; +} + +type Listener = () => void; + +/** + * Module-level singleton external store for chat streaming. + * + * Single-active-stream invariant: at most one turn streams at a time, + * tracked by {@link active}. Per-thread state is still keyed by threadId so + * the currently-streaming thread and the currently-viewed thread can differ + * (e.g. a stream finishes while the user is on another chat). + * + * ponytail: unbounded ``states`` map is bounded in practice by + * ``clearInactive`` (called on navigation) + ``clear`` (called after the DB + * re-hydrates a finished turn). Upgrade path if that ever leaks: LRU cap. + */ +class ChatStreamStore { + private states = new Map(); + private listeners = new Set(); + + /** Shared, cross-navigation token-usage store (one instance app-wide). */ + readonly tokenUsage = createTokenUsageStore(); + + /** The one in-flight turn's abort handle, or null when idle. */ + private active: { threadId: number; controller: AbortController } | null = null; + + /** Timestamp of the last explicit cancel, for the THREAD_BUSY retry window. */ + recentCancelRequestedAt = 0; + + subscribe = (listener: Listener): (() => void) => { + this.listeners.add(listener); + return () => { + this.listeners.delete(listener); + }; + }; + + private notify(): void { + for (const l of this.listeners) l(); + } + + /** Snapshot for ``useSyncExternalStore``; stable ref between mutations. */ + getSnapshot = (threadId: number | null): ThreadStreamState | null => { + if (threadId == null) return null; + return this.states.get(threadId) ?? null; + }; + + isRunning(threadId: number | null): boolean { + if (threadId == null) return false; + return this.states.get(threadId)?.isRunning ?? false; + } + + getMessages(threadId: number): ThreadMessageLike[] { + return this.states.get(threadId)?.messages ?? []; + } + + getPendingInterrupts(threadId: number): PendingInterruptState[] { + return this.states.get(threadId)?.pendingInterrupts ?? []; + } + + private ensure(threadId: number): ThreadStreamState { + let s = this.states.get(threadId); + if (!s) { + s = { threadId, messages: [], isRunning: false, pendingInterrupts: [] }; + this.states.set(threadId, s); + } + return s; + } + + private commit(threadId: number, next: ThreadStreamState): void { + this.states.set(threadId, next); + this.notify(); + } + + /** Seed a thread's state at the start of a fresh turn (running=true). */ + begin(threadId: number, messages: ThreadMessageLike[]): void { + this.commit(threadId, { threadId, messages, isRunning: true, pendingInterrupts: [] }); + } + + setMessages(threadId: number, updater: (prev: ThreadMessageLike[]) => ThreadMessageLike[]): void { + const prev = this.ensure(threadId); + const messages = updater(prev.messages); + if (messages === prev.messages) return; + this.commit(threadId, { ...prev, messages }); + } + + setRunning(threadId: number, running: boolean): void { + const prev = this.ensure(threadId); + if (prev.isRunning === running) return; + this.commit(threadId, { ...prev, isRunning: running }); + } + + setPendingInterrupts( + threadId: number, + updater: (prev: PendingInterruptState[]) => PendingInterruptState[] + ): void { + const prev = this.ensure(threadId); + const pendingInterrupts = updater(prev.pendingInterrupts); + if (pendingInterrupts === prev.pendingInterrupts) return; + this.commit(threadId, { ...prev, pendingInterrupts }); + } + + /** + * A thread whose overlay must survive DB re-hydration / navigation: it is + * either streaming or paused awaiting a HITL decision (the pending + * interrupts + interrupt cards live only in the overlay). + */ + private isPinned(s: ThreadStreamState): boolean { + return s.isRunning || s.pendingInterrupts.length > 0; + } + + /** Drop a thread's overlay once the DB is authoritative. No-op while pinned. */ + clear(threadId: number): void { + const s = this.states.get(threadId); + if (!s || this.isPinned(s)) return; + this.states.delete(threadId); + this.notify(); + } + + /** Evict every non-pinned thread except ``exceptThreadId`` (memory bound). */ + clearInactive(exceptThreadId: number | null): void { + let changed = false; + for (const [id, s] of this.states) { + if (id === exceptThreadId || this.isPinned(s)) continue; + this.states.delete(id); + changed = true; + } + if (changed) this.notify(); + } + + // ---- active-stream lifecycle ------------------------------------------- + + /** Register a new in-flight turn, aborting any previous one first. */ + beginActive(threadId: number, controller: AbortController): void { + this.abortActive(); + this.active = { threadId, controller }; + } + + get activeThreadId(): number | null { + return this.active?.threadId ?? null; + } + + /** Clear the active handle iff it still points at ``controller``. */ + clearActive(controller: AbortController): void { + if (this.active?.controller === controller) this.active = null; + } + + /** Abort the in-flight turn's fetch (client disconnect, not a server stop). */ + abortActive(): void { + if (this.active) { + this.active.controller.abort(); + this.active = null; + } + } + + markRecentCancel(): void { + this.recentCancelRequestedAt = Date.now(); + } +} + +export const chatStreamStore = new ChatStreamStore(); diff --git a/surfsense_web/lib/chat/stream-engine/use-chat-stream.ts b/surfsense_web/lib/chat/stream-engine/use-chat-stream.ts new file mode 100644 index 000000000..ccfb32481 --- /dev/null +++ b/surfsense_web/lib/chat/stream-engine/use-chat-stream.ts @@ -0,0 +1,18 @@ +"use client"; + +import { useCallback, useSyncExternalStore } from "react"; +import { chatStreamStore, type ThreadStreamState } from "./store"; + +/** + * Subscribe to the durable streaming state for a given thread. + * + * Returns the live ``ThreadStreamState`` while a turn is streaming (or + * pending re-hydration from the DB), or ``null`` when the thread has no + * active overlay — in which case the page falls back to its DB-hydrated + * messages. State lives in the module-level {@link chatStreamStore}, so it + * survives the chat page unmounting during in-app navigation. + */ +export function useChatStream(threadId: number | null): ThreadStreamState | null { + const getSnapshot = useCallback(() => chatStreamStore.getSnapshot(threadId), [threadId]); + return useSyncExternalStore(chatStreamStore.subscribe, getSnapshot, () => null); +} diff --git a/surfsense_web/lib/format-date.ts b/surfsense_web/lib/format-date.ts index c2f445537..a062b08fa 100644 --- a/surfsense_web/lib/format-date.ts +++ b/surfsense_web/lib/format-date.ts @@ -68,3 +68,17 @@ export function formatRelativeFutureDate(dateString: string): string { export function formatThreadTimestamp(dateString: string): string { return format(new Date(dateString), "MMM d, yyyy 'at' h:mm a"); } + +/** + * Format a chat message timestamp for inline display under a bubble. + * Locale-aware, 12-hour clock. Example: "Jul 13, 10:42 PM". + */ +export function formatMessageTimestamp(date: Date): string { + return date.toLocaleDateString(undefined, { + month: "short", + day: "numeric", + hour: "numeric", + minute: "2-digit", + hour12: true, + }); +}