mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-10-08 10:47:48 +08:00
dayuan.jiang
366d9ee4ff
The canvas now fills the window and the chat floats on its right. draw.io is served from our own origin (public/drawio, downloaded at build time and checked against the release's SHA-256), so the app can drive the editor directly: - every AI change becomes a version card with a thumbnail; compare, restore or undo it, and Ctrl+Z on the canvas takes an AI change back in one step - shapes the AI just changed are highlighted on the canvas - selected shapes appear as a chip and are sent to the model - the app's own canvas toolbar, shape library and page tabs replace draw.io's chrome - new start screen, session rename and switcher, open .drawio files - settings split into Models & keys, Appearance, Advanced and About - phones switch between canvas and chat - dark mode keeps the diagram on a white sheet and no longer flashes With an external draw.io (NEXT_PUBLIC_DRAWIO_BASE_URL) the app falls back to draw.io's own toolbar and the features above that need same origin are off. The chat logic moved from the 1500-line chat panel into a chat engine plus small UI components; zustand holds canvas, versions, settings and UI state.
1504 lines
56 KiB
TypeScript
1504 lines
56 KiB
TypeScript
"use client"
|
|
|
|
import { useChat } from "@ai-sdk/react"
|
|
import {
|
|
type ChatStatus,
|
|
DefaultChatTransport,
|
|
isToolUIPart,
|
|
type UIMessage,
|
|
} from "ai"
|
|
import { usePathname, useRouter, useSearchParams } from "next/navigation"
|
|
import type React from "react"
|
|
import {
|
|
createContext,
|
|
useCallback,
|
|
useContext,
|
|
useEffect,
|
|
useLayoutEffect,
|
|
useMemo,
|
|
useRef,
|
|
useState,
|
|
} from "react"
|
|
import { flushSync } from "react-dom"
|
|
import { toast } from "sonner"
|
|
import type { ValidationState } from "@/components/chat/ValidationCard"
|
|
import { useDiagram } from "@/contexts/diagram-context"
|
|
import { useDiagramToolHandlers } from "@/hooks/use-diagram-tool-handlers"
|
|
import { useDictionary } from "@/hooks/use-dictionary"
|
|
import {
|
|
getSelectedAIConfig,
|
|
type UseModelConfigReturn,
|
|
useModelConfig,
|
|
} from "@/hooks/use-model-config"
|
|
import {
|
|
getSessionVersions,
|
|
type SessionData,
|
|
useSessionManager,
|
|
} from "@/hooks/use-session-manager"
|
|
import { useValidateDiagram } from "@/hooks/use-validate-diagram"
|
|
import { getApiEndpoint } from "@/lib/base-path"
|
|
import { findCachedResponse } from "@/lib/cached-responses"
|
|
import { buildChatHeaders } from "@/lib/chat-request"
|
|
import { formatMessage } from "@/lib/i18n/utils"
|
|
import { isPdfFile, isTextFile } from "@/lib/pdf-utils"
|
|
import type { SessionMetadata } from "@/lib/session-storage"
|
|
import { sanitizeMessages } from "@/lib/session-storage"
|
|
import { STORAGE_KEYS } from "@/lib/storage"
|
|
import type { UrlData } from "@/lib/url-utils"
|
|
import { type FileData, useFileProcessor } from "@/lib/use-file-processor"
|
|
import { useQuotaManager } from "@/lib/use-quota-manager"
|
|
import { formatXML, isRealDiagram } from "@/lib/utils"
|
|
import { parseDrawioFileContent } from "@/packages/mcp-server/src/load-diagram.ts"
|
|
import { prepareNewDiagram } from "@/packages/mcp-server/src/new-diagram.ts"
|
|
import { BLANK_MXFILE, hasCells } from "@/packages/mcp-server/src/pages.ts"
|
|
import { type SelectedCell, useCanvasStore } from "@/stores/canvas-store"
|
|
import { useSettingsStore } from "@/stores/settings-store"
|
|
import { useUiStore } from "@/stores/ui-store"
|
|
import { useVersionsStore } from "@/stores/versions-store"
|
|
import {
|
|
APPENDED_FILE_SECTIONS_PATTERN,
|
|
getUserOriginalText,
|
|
} from "./message-text"
|
|
import { useDiagramStreaming } from "./use-diagram-streaming"
|
|
import { useVersions } from "./use-versions"
|
|
|
|
// sessionStorage key for the unsent input (survives remounts and reloads)
|
|
const SESSION_STORAGE_INPUT_KEY = "next-ai-draw-io-input"
|
|
|
|
const TOOL_ERROR_STATE = "output-error" as const
|
|
// 3 to support VLM validation retries (matches MAX_VALIDATION_RETRIES)
|
|
const MAX_AUTO_RETRY_COUNT = 3
|
|
// Limit for truncation continuation retries
|
|
const MAX_CONTINUATION_RETRY_COUNT = 2
|
|
// Save the session at most once per second
|
|
const SAVE_DEBOUNCE_MS = 1000
|
|
|
|
/**
|
|
* Check if auto-resubmit should happen based on tool errors.
|
|
* Only checks the LAST tool part (most recent tool call), not all tool parts.
|
|
*/
|
|
function hasToolErrors(messages: UIMessage[]): boolean {
|
|
const lastMessage = messages[messages.length - 1]
|
|
if (lastMessage?.role !== "assistant") return false
|
|
const lastToolPart = lastMessage.parts.filter(isToolUIPart).at(-1)
|
|
return lastToolPart?.state === TOOL_ERROR_STATE
|
|
}
|
|
|
|
/**
|
|
* Snapshots keep the full multi-page document, but the model only sees and
|
|
* edits the first page, so give it the first page's mxGraphModel.
|
|
* Older snapshots already hold a single mxGraphModel and are returned as is.
|
|
*/
|
|
function getFirstPageXml(xml: string): string {
|
|
if (!xml.includes("<mxfile")) return xml
|
|
const doc = new DOMParser().parseFromString(xml, "text/xml")
|
|
const model = doc.querySelector("diagram")?.querySelector("mxGraphModel")
|
|
return model ? formatXML(new XMLSerializer().serializeToString(model)) : xml
|
|
}
|
|
|
|
// Shapes sent with a user message (also kept in its metadata)
|
|
type MessageSelection = Pick<SelectedCell, "id" | "label">[]
|
|
|
|
function selectionOf(message: UIMessage): MessageSelection {
|
|
const meta = message.metadata as
|
|
| { selectedCells?: MessageSelection }
|
|
| undefined
|
|
return Array.isArray(meta?.selectedCells) ? meta.selectedCells : []
|
|
}
|
|
|
|
function newLangfuseSessionId() {
|
|
return `session-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`
|
|
}
|
|
|
|
/**
|
|
* Read the diagram document out of a .drawio, .xml or .drawio.svg file, with
|
|
* its pages inflated: previews and the model read it before draw.io saves
|
|
* it again. A broken or cut off file gives null.
|
|
*/
|
|
async function readDiagramFile(file: File): Promise<string | null> {
|
|
let text = await file.text()
|
|
// Editable SVG: the diagram is stored in the root's `content` attribute
|
|
if (text.includes("<svg")) {
|
|
const doc = new DOMParser().parseFromString(text, "image/svg+xml")
|
|
text = doc.documentElement?.getAttribute("content") ?? ""
|
|
}
|
|
const loaded = parseDrawioFileContent(text)
|
|
return loaded.ok ? loaded.xml : null
|
|
}
|
|
|
|
export interface ChatEngine {
|
|
messages: UIMessage[]
|
|
status: ChatStatus
|
|
error: Error | undefined
|
|
isBusy: boolean
|
|
/** The chat on screen is being left (saved, then another one shown) */
|
|
isLeaving: boolean
|
|
/** A message is being sent again; the chat is not empty meanwhile */
|
|
isResending: boolean
|
|
setMessages: (messages: UIMessage[]) => void
|
|
stop: () => void
|
|
|
|
input: string
|
|
setInput: (value: string) => void
|
|
files: File[]
|
|
pdfData: Map<File, FileData>
|
|
setFiles: (files: File[]) => void
|
|
urlData: Map<string, UrlData>
|
|
setUrlData: React.Dispatch<React.SetStateAction<Map<string, UrlData>>>
|
|
isExtractingAttachments: boolean
|
|
|
|
submit: (overrideText?: string) => Promise<void>
|
|
sendTemplate: (prompt: string) => void
|
|
improveWithSuggestions: (feedback: string) => Promise<void>
|
|
regenerate: (messageIndex: number) => Promise<void>
|
|
retryLastMessage: () => void
|
|
editMessage: (messageIndex: number, newText: string) => Promise<void>
|
|
|
|
/** Selected shapes that will be sent with the next message */
|
|
chatSelection: SelectedCell[]
|
|
dismissSelection: () => void
|
|
|
|
sessions: SessionMetadata[]
|
|
currentSessionId: string | null
|
|
currentTitle: string | null
|
|
isRestored: boolean
|
|
newChat: () => Promise<void>
|
|
selectSession: (id: string) => Promise<void>
|
|
deleteSession: (id: string) => Promise<void>
|
|
renameSession: (title: string) => Promise<void>
|
|
openDiagramFile: (file: File) => Promise<void>
|
|
/**
|
|
* Save the chat on screen, then leave the page (another language) with
|
|
* the chat's session id; nothing happens while an answer runs
|
|
*/
|
|
leavePage: (go: (sessionId: string | null) => void) => Promise<void>
|
|
|
|
restoreVersion: (versionId: string) => void
|
|
undoVersion: (versionId: string) => void
|
|
|
|
langfuseSessionId: string
|
|
loadedMessageIdsRef: React.RefObject<Set<string>>
|
|
validationStates: Record<string, ValidationState>
|
|
modelConfig: UseModelConfigReturn
|
|
}
|
|
|
|
const ChatEngineContext = createContext<ChatEngine | null>(null)
|
|
|
|
export function useChatEngine(): ChatEngine {
|
|
const engine = useContext(ChatEngineContext)
|
|
if (!engine) {
|
|
throw new Error("useChatEngine must be used inside ChatEngineProvider")
|
|
}
|
|
return engine
|
|
}
|
|
|
|
export function ChatEngineProvider({
|
|
children,
|
|
}: {
|
|
children: React.ReactNode
|
|
}) {
|
|
const {
|
|
loadDiagram: onDisplayChart,
|
|
handleExport: onExport,
|
|
exportResolversRef,
|
|
chartXML,
|
|
chartXMLRef: liveChartXMLRef,
|
|
latestSvg,
|
|
clearDiagram,
|
|
getThumbnailSvg,
|
|
captureValidationPng,
|
|
startTurn,
|
|
} = useDiagram()
|
|
|
|
const dict = useDictionary()
|
|
const router = useRouter()
|
|
const pathname = usePathname()
|
|
const searchParams = useSearchParams()
|
|
const urlSessionId = searchParams.get("session")
|
|
|
|
const settings = useSettingsStore()
|
|
const openSettings = useUiStore((s) => s.openSettings)
|
|
const focusComposer = useUiStore((s) => s.focusComposer)
|
|
|
|
const onFetchChart = () => {
|
|
// Waits for the reply to its own export, by its tag
|
|
const tag = onExport()
|
|
return Promise.race([
|
|
new Promise<string>((resolve) => {
|
|
if (tag) exportResolversRef.current[tag] = resolve
|
|
}),
|
|
new Promise<string>((_, reject) => {
|
|
setTimeout(() => {
|
|
delete exportResolversRef.current[tag]
|
|
reject(new Error("Chart export timed out after 10 seconds"))
|
|
}, 10000)
|
|
}),
|
|
])
|
|
}
|
|
|
|
const { files, pdfData, handleFileChange, setFiles } = useFileProcessor({
|
|
tooLong: dict.errors.fileTooLong,
|
|
readFailed: dict.errors.fileReadFailed,
|
|
})
|
|
const [urlData, setUrlData] = useState<Map<string, UrlData>>(new Map())
|
|
const filesRef = useRef(files)
|
|
filesRef.current = files
|
|
|
|
const modelConfig = useModelConfig()
|
|
const sessionManager = useSessionManager({ initialSessionId: urlSessionId })
|
|
|
|
const [input, setInputState] = useState("")
|
|
const [dailyRequestLimit, setDailyRequestLimit] = useState(0)
|
|
const [dailyTokenLimit, setDailyTokenLimit] = useState(0)
|
|
const [tpmLimit, setTpmLimit] = useState(0)
|
|
|
|
// Restore the unsent input from sessionStorage
|
|
useEffect(() => {
|
|
const savedInput = sessionStorage.getItem(SESSION_STORAGE_INPUT_KEY)
|
|
if (savedInput) {
|
|
inputRef.current = savedInput
|
|
setInputState(savedInput)
|
|
}
|
|
}, [])
|
|
|
|
// Latest composer contents, for code that runs after an await
|
|
const inputRef = useRef("")
|
|
const setInput = useCallback((value: string) => {
|
|
inputRef.current = value
|
|
setInputState(value)
|
|
if (value) sessionStorage.setItem(SESSION_STORAGE_INPUT_KEY, value)
|
|
else sessionStorage.removeItem(SESSION_STORAGE_INPUT_KEY)
|
|
}, [])
|
|
|
|
useEffect(() => {
|
|
fetch(getApiEndpoint("/api/config"))
|
|
.then((res) => res.json())
|
|
.then((data) => {
|
|
setDailyRequestLimit(data.dailyRequestLimit || 0)
|
|
setDailyTokenLimit(data.dailyTokenLimit || 0)
|
|
setTpmLimit(data.tpmLimit || 0)
|
|
})
|
|
.catch(() => {})
|
|
}, [])
|
|
|
|
const quotaManager = useQuotaManager({
|
|
dailyRequestLimit,
|
|
dailyTokenLimit,
|
|
tpmLimit,
|
|
onConfigModel: () => openSettings("models"),
|
|
})
|
|
|
|
// Session ID for Langfuse tracing (restored from localStorage if available)
|
|
const [langfuseSessionId, setLangfuseSessionId] = useState(() => {
|
|
if (typeof window !== "undefined") {
|
|
const saved = localStorage.getItem(STORAGE_KEYS.sessionId)
|
|
if (saved) return saved
|
|
}
|
|
return newLangfuseSessionId()
|
|
})
|
|
useEffect(() => {
|
|
localStorage.setItem(STORAGE_KEYS.sessionId, langfuseSessionId)
|
|
}, [langfuseSessionId])
|
|
|
|
// Title for a session that is not saved yet (opened file, early rename)
|
|
const [pendingTitle, setPendingTitleState] = useState<string | null>(null)
|
|
const pendingTitleRef = useRef<string | null>(null)
|
|
const setPendingTitle = (title: string | null) => {
|
|
pendingTitleRef.current = title
|
|
setPendingTitleState(title)
|
|
}
|
|
|
|
// XML snapshot taken before each user message (keyed by message index)
|
|
const xmlSnapshotsRef = useRef<Map<number, string>>(new Map())
|
|
// Index of the user message whose turn is running (versions belong to it)
|
|
const currentTurnRef = useRef(0)
|
|
|
|
const hasRestoredRef = useRef(false)
|
|
const [isRestored, setIsRestored] = useState(false)
|
|
|
|
// Latest chartXML for callbacks (avoids stale closures)
|
|
const chartXMLRef = useRef(chartXML)
|
|
// Session loaded without a diagram (prevents thumbnail contamination)
|
|
const justLoadedSessionIdRef = useRef<string | null>(null)
|
|
useEffect(() => {
|
|
chartXMLRef.current = chartXML
|
|
if (chartXML) justLoadedSessionIdRef.current = null
|
|
}, [chartXML])
|
|
|
|
const latestSvgRef = useRef(latestSvg)
|
|
useEffect(() => {
|
|
latestSvgRef.current = latestSvg
|
|
}, [latestSvg])
|
|
|
|
const autoRetryCountRef = useRef(0)
|
|
const continuationRetryCountRef = useRef(0)
|
|
// Partial XML when output was truncated; non-empty means continuation mode
|
|
const partialXmlRef = useRef<string>("")
|
|
// Diagram before a cut off drawing, whose half drawn preview stays on
|
|
// the canvas while append_diagram finishes it
|
|
const continuationOriginalRef = useRef<string | null>(null)
|
|
// Tool calls already applied, so remounts never replay old outputs
|
|
const processedToolCallsRef = useRef<Set<string>>(new Set())
|
|
// Set by Stop until the user sends the next message
|
|
const stoppedRef = useRef(false)
|
|
const preparingSendRef = useRef(false)
|
|
// Set while the chat on screen is left (saved, then another one put on
|
|
// screen): an answer started meanwhile would land in the next chat
|
|
const leavingRef = useRef(false)
|
|
const [isLeaving, setIsLeaving] = useState(false)
|
|
// A message is being sent again (retry, regenerate, edit), until the
|
|
// SDK has put it back
|
|
const [isResending, setIsResending] = useState(false)
|
|
// Presses of Stop: a check that began before one still knows of it after
|
|
// the next message clears stoppedRef
|
|
const stopCountRef = useRef(0)
|
|
// Diagram before each streamed display_diagram and edit_diagram preview
|
|
// (key: toolCallId), shared between the streaming preview and the tool
|
|
// handler
|
|
const editDiagramOriginalXmlRef = useRef<Map<string, string>>(new Map())
|
|
const saveDebounceRef = useRef<ReturnType<typeof setTimeout> | null>(null)
|
|
|
|
const [validationStates, setValidationStates] = useState<
|
|
Record<string, ValidationState>
|
|
>({})
|
|
const handleValidationStateChange = useCallback(
|
|
(toolCallId: string, state: ValidationState) => {
|
|
setValidationStates((prev) => ({ ...prev, [toolCallId]: state }))
|
|
},
|
|
[],
|
|
)
|
|
|
|
// Failed VLM validations in the current user turn (reset on user action)
|
|
const validationRetryCountRef = useRef(0)
|
|
const { validateWithFallback, cancel: cancelValidation } =
|
|
useValidateDiagram()
|
|
|
|
const { handleToolCall } = useDiagramToolHandlers({
|
|
partialXmlRef,
|
|
continuationOriginalRef,
|
|
editDiagramOriginalXmlRef,
|
|
processedToolCallsRef,
|
|
validationRetryCountRef,
|
|
// A preview undone just before the tool call is in this one already
|
|
chartXMLRef: liveChartXMLRef,
|
|
onDisplayChart,
|
|
onFetchChart,
|
|
captureValidationPng,
|
|
validateDiagram: validateWithFallback,
|
|
enableVlmValidation: settings.vlmValidationEnabled,
|
|
sessionId: langfuseSessionId,
|
|
watchStop: () => {
|
|
const stopsBefore = stopCountRef.current
|
|
return () =>
|
|
stoppedRef.current || stopCountRef.current !== stopsBefore
|
|
},
|
|
onValidationStateChange: handleValidationStateChange,
|
|
})
|
|
|
|
// The turn ends without finishing a cut off drawing: back to the
|
|
// diagram before it
|
|
const abandonContinuation = () => {
|
|
const original = continuationOriginalRef.current
|
|
if (original === null) return
|
|
continuationOriginalRef.current = null
|
|
// Previews drawn since are on top of the cut off one: none of them
|
|
// may be undone to it later
|
|
for (const id of editDiagramOriginalXmlRef.current.keys()) {
|
|
processedToolCallsRef.current.add(id)
|
|
}
|
|
editDiagramOriginalXmlRef.current.clear()
|
|
onDisplayChart(original, true, "revert")
|
|
}
|
|
|
|
const transport = useMemo(
|
|
() => new DefaultChatTransport({ api: getApiEndpoint("/api/chat") }),
|
|
[],
|
|
)
|
|
|
|
const {
|
|
messages,
|
|
sendMessage,
|
|
addToolOutput,
|
|
status,
|
|
error,
|
|
setMessages,
|
|
stop,
|
|
} = useChat({
|
|
transport,
|
|
onToolCall: async ({ toolCall }) => {
|
|
await handleToolCall({ toolCall }, addToolOutput)
|
|
},
|
|
onFinish: ({ message, isAbort, isError }) => {
|
|
// Stopped or failed: tool calls still streaming never reach the
|
|
// tool handler. Mark them handled so a later render of the
|
|
// stream does not draw their preview again.
|
|
if (!isAbort && !isError) return
|
|
for (const part of message.parts as any[]) {
|
|
if (part.state === "input-streaming" && part.toolCallId) {
|
|
processedToolCallsRef.current.add(part.toolCallId)
|
|
}
|
|
}
|
|
},
|
|
onError: (error) => {
|
|
// A diagram still streaming when the request failed never
|
|
// reaches the tool handler: undo its preview. Only previews not
|
|
// handled yet are stored, and the first one holds the diagram
|
|
// before any of them.
|
|
const [originalXml] = editDiagramOriginalXmlRef.current.values()
|
|
if (originalXml) onDisplayChart(originalXml, true, "revert")
|
|
editDiagramOriginalXmlRef.current.clear()
|
|
abandonContinuation()
|
|
|
|
// Server errors are JSON: a quota limit ({type: request, token or
|
|
// tpm}), a provider error ({type: "provider", code, message}) or
|
|
// {error}. The SDK puts the response body in error.message.
|
|
let data: any = null
|
|
try {
|
|
data = JSON.parse(error.message)
|
|
} catch {
|
|
// Plain text, e.g. a network failure in the browser
|
|
}
|
|
if (data?.type === "request") {
|
|
quotaManager.showQuotaLimitToast(data.used, data.limit)
|
|
return
|
|
}
|
|
if (data?.type === "token") {
|
|
quotaManager.showTokenLimitToast(data.used, data.limit)
|
|
return
|
|
}
|
|
if (data?.type === "tpm") {
|
|
quotaManager.showTPMLimitToast(data.limit)
|
|
return
|
|
}
|
|
|
|
const isAccessCodeError = String(
|
|
data?.error ?? error.message,
|
|
).includes("Invalid or missing access code")
|
|
if (!isAccessCodeError) console.error("Chat error:", error)
|
|
|
|
// A hint the user can act on, then the provider's own words
|
|
let text: string = error.message
|
|
let openModelConfig = false
|
|
if (data?.type === "provider") {
|
|
const hints = dict.errors.llm as Record<string, string>
|
|
const hint = hints[data.code]
|
|
text =
|
|
hint && data.message
|
|
? `${hint}\n\n${data.message}`
|
|
: hint || data.message
|
|
openModelConfig = [
|
|
"invalid_api_key",
|
|
"forbidden",
|
|
"model_not_found",
|
|
"server_key_forbidden",
|
|
].includes(data.code)
|
|
} else if (typeof data?.error === "string") {
|
|
text = data.error
|
|
} else if (error.message === "Failed to fetch") {
|
|
text = dict.errors.networkError
|
|
}
|
|
|
|
// A system message, so it can be cleared with the conversation
|
|
setMessages((currentMessages) => [
|
|
...currentMessages,
|
|
{
|
|
id: `error-${Date.now()}`,
|
|
role: "system" as const,
|
|
parts: [{ type: "text" as const, text }],
|
|
// The message shows a button that opens model settings
|
|
...(openModelConfig && {
|
|
metadata: { openModelConfig: true },
|
|
}),
|
|
},
|
|
])
|
|
|
|
if (isAccessCodeError) openSettings("advanced")
|
|
},
|
|
// Re-render streamed messages at most every 150 ms. The streaming
|
|
// diagram preview draws on each update, so this also limits redraws
|
|
experimental_throttle: 150,
|
|
sendAutomaticallyWhen: ({ messages }) => {
|
|
const send = shouldSendAgain(messages)
|
|
// The turn ends here
|
|
if (!send) abandonContinuation()
|
|
return send
|
|
},
|
|
})
|
|
|
|
// After each request: send the next one (retry, continue a cut off
|
|
// drawing) or end the turn
|
|
const shouldSendAgain = (messages: UIMessage[]) => {
|
|
// The user stopped: a tool result that arrives later (a VLM
|
|
// check still running) must not start a new request
|
|
if (stoppedRef.current) return false
|
|
|
|
const isInContinuationMode = partialXmlRef.current.length > 0
|
|
const shouldRetry = hasToolErrors(messages)
|
|
|
|
if (!shouldRetry) {
|
|
autoRetryCountRef.current = 0
|
|
continuationRetryCountRef.current = 0
|
|
partialXmlRef.current = ""
|
|
return false
|
|
}
|
|
|
|
if (isInContinuationMode) {
|
|
if (
|
|
continuationRetryCountRef.current >=
|
|
MAX_CONTINUATION_RETRY_COUNT
|
|
) {
|
|
toast.error(
|
|
formatMessage(dict.errors.continuationRetryLimit, {
|
|
max: MAX_CONTINUATION_RETRY_COUNT,
|
|
}),
|
|
)
|
|
continuationRetryCountRef.current = 0
|
|
partialXmlRef.current = ""
|
|
return false
|
|
}
|
|
continuationRetryCountRef.current++
|
|
} else {
|
|
if (autoRetryCountRef.current >= MAX_AUTO_RETRY_COUNT) {
|
|
toast.error(
|
|
formatMessage(dict.errors.retryLimit, {
|
|
max: MAX_AUTO_RETRY_COUNT,
|
|
}),
|
|
)
|
|
autoRetryCountRef.current = 0
|
|
partialXmlRef.current = ""
|
|
return false
|
|
}
|
|
autoRetryCountRef.current++
|
|
}
|
|
return true
|
|
}
|
|
|
|
const isBusy = status === "streaming" || status === "submitted"
|
|
useEffect(() => {
|
|
if (messages.length > 0) setIsResending(false)
|
|
}, [messages.length])
|
|
// isBusy of the latest render, for code that runs after an await
|
|
const busyRef = useRef(isBusy)
|
|
busyRef.current = isBusy
|
|
// The page goes (another language mounts a new one): the answer stops,
|
|
// so it cannot reach the next page's canvas
|
|
const stopRef = useRef(stop)
|
|
stopRef.current = stop
|
|
useEffect(
|
|
() => () => {
|
|
stopRef.current()
|
|
},
|
|
[],
|
|
)
|
|
// While an answer is sent or streams in, the chat must stay on screen:
|
|
// the rest of the answer and its diagram would land in another chat
|
|
// (Between the requests of one turn the SDK is idle for a moment; a cut
|
|
// off drawing still waiting for its continuation keeps the turn going)
|
|
const isTurnActive = () =>
|
|
busyRef.current ||
|
|
preparingSendRef.current ||
|
|
continuationOriginalRef.current !== null
|
|
// Sending, version actions: not during an answer, nor while leaving
|
|
const canStartTurn = () => !isTurnActive() && !leavingRef.current
|
|
|
|
// Message IDs loaded from a session (skip animations and replays)
|
|
const loadedMessageIdsRef = useRef<Set<string>>(new Set())
|
|
|
|
useDiagramStreaming({
|
|
messages,
|
|
processedToolCallsRef,
|
|
editDiagramOriginalXmlRef,
|
|
loadedMessageIdsRef,
|
|
})
|
|
|
|
const { restoreVersion, undoVersion } = useVersions({
|
|
currentTurnRef,
|
|
isBusy,
|
|
isTurnActive: () => !canStartTurn(),
|
|
getChatGeneration: sessionManager.getChatGeneration,
|
|
})
|
|
|
|
// ---------------------------------------------------------------------
|
|
// Selection sent with the next message
|
|
// ---------------------------------------------------------------------
|
|
|
|
const selection = useCanvasStore((s) => s.selection)
|
|
// The model sees and edits the first page only
|
|
const onFirstPage = useCanvasStore(
|
|
(s) => s.pages.length === 0 || s.pages[0]?.id === s.currentPageId,
|
|
)
|
|
const [dismissedSelectionKey, setDismissedSelectionKey] = useState("")
|
|
const selectionKey = selection.map((c) => c.id).join(",")
|
|
const chatSelection =
|
|
!onFirstPage || selectionKey === dismissedSelectionKey ? [] : selection
|
|
// Once nothing is selected, picking the same shapes again attaches them
|
|
useEffect(() => {
|
|
if (selection.length === 0) setDismissedSelectionKey("")
|
|
}, [selection.length])
|
|
const dismissSelection = useCallback(
|
|
() => setDismissedSelectionKey(selectionKey),
|
|
[selectionKey],
|
|
)
|
|
|
|
// ---------------------------------------------------------------------
|
|
// Sessions
|
|
// ---------------------------------------------------------------------
|
|
|
|
const messagesRef = useRef(messages)
|
|
useEffect(() => {
|
|
messagesRef.current = messages
|
|
}, [messages])
|
|
|
|
// Last synced session ID, to detect external changes (URL back/forward)
|
|
const lastSyncedSessionIdRef = useRef<string | null>(null)
|
|
// Message arrays of our own saves. A session holding one of them was
|
|
// created by our own save, so it must not be treated as an external
|
|
// switch (with two saves of a new chat at once, the first creates it).
|
|
const savedMessagesRef = useRef(new WeakSet<object>())
|
|
// Session was just loaded (skip the auto-save that would reorder it)
|
|
const justLoadedSessionRef = useRef(false)
|
|
|
|
const syncUIWithSession = useCallback(
|
|
(
|
|
data:
|
|
| (SessionData & {
|
|
diagramHistory?: { svg: string; xml: string }[]
|
|
})
|
|
| null,
|
|
) => {
|
|
const hasRealDiagram = isRealDiagram(data?.diagramXml)
|
|
// They act on the chat that was on screen: its canvas (Undo of
|
|
// a restore) or leaving it unsaved
|
|
toast.dismiss("version-restored")
|
|
toast.dismiss("session-save-leave")
|
|
if (data) {
|
|
loadedMessageIdsRef.current = new Set(
|
|
(data.messages as { id: string }[]).map((m) => m.id),
|
|
)
|
|
setMessages(data.messages as unknown as UIMessage[])
|
|
xmlSnapshotsRef.current = new Map(data.xmlSnapshots)
|
|
if (hasRealDiagram) {
|
|
onDisplayChart(data.diagramXml, true)
|
|
chartXMLRef.current = data.diagramXml
|
|
} else {
|
|
clearDiagram()
|
|
chartXMLRef.current = ""
|
|
latestSvgRef.current = ""
|
|
}
|
|
useVersionsStore
|
|
.getState()
|
|
.setVersions(getSessionVersions(data))
|
|
} else {
|
|
loadedMessageIdsRef.current = new Set()
|
|
setMessages([])
|
|
xmlSnapshotsRef.current.clear()
|
|
clearDiagram()
|
|
chartXMLRef.current = ""
|
|
latestSvgRef.current = ""
|
|
useVersionsStore.getState().clear()
|
|
}
|
|
},
|
|
[setMessages, onDisplayChart, clearDiagram],
|
|
)
|
|
|
|
const buildSessionData = useCallback(
|
|
async (
|
|
options: { withThumbnail?: boolean } = {},
|
|
): Promise<SessionData> => {
|
|
// Read together, before the wait for the thumbnail: what changes
|
|
// meanwhile goes into the next save, not half into this one
|
|
const currentDiagramXml = chartXMLRef.current || ""
|
|
const messages = sanitizeMessages(messagesRef.current)
|
|
savedMessagesRef.current.add(messages)
|
|
const data = {
|
|
messages,
|
|
xmlSnapshots: Array.from(xmlSnapshotsRef.current.entries()),
|
|
diagramXml: currentDiagramXml,
|
|
versions: useVersionsStore.getState().versions,
|
|
title: pendingTitleRef.current ?? undefined,
|
|
}
|
|
let thumbnailDataUrl: string | undefined
|
|
if (isRealDiagram(currentDiagramXml) && options.withThumbnail) {
|
|
const freshThumb = await getThumbnailSvg()
|
|
if (freshThumb) {
|
|
latestSvgRef.current = freshThumb
|
|
thumbnailDataUrl = freshThumb
|
|
} else if (latestSvgRef.current) {
|
|
thumbnailDataUrl = latestSvgRef.current
|
|
}
|
|
}
|
|
return { ...data, thumbnailDataUrl }
|
|
},
|
|
[getThumbnailSvg],
|
|
)
|
|
|
|
// Restore from the session manager once it has loaded
|
|
useLayoutEffect(() => {
|
|
if (hasRestoredRef.current) return
|
|
if (sessionManager.isLoading) return
|
|
hasRestoredRef.current = true
|
|
try {
|
|
const currentSession = sessionManager.currentSession
|
|
if (currentSession) {
|
|
justLoadedSessionRef.current = true
|
|
syncUIWithSession(currentSession)
|
|
}
|
|
lastSyncedSessionIdRef.current = sessionManager.currentSessionId
|
|
} catch (error) {
|
|
console.error("Failed to restore session:", error)
|
|
toast.error(dict.errors.sessionCorrupted)
|
|
} finally {
|
|
setIsRestored(true)
|
|
}
|
|
}, [
|
|
sessionManager.isLoading,
|
|
sessionManager.currentSession,
|
|
syncUIWithSession,
|
|
dict.errors.sessionCorrupted,
|
|
])
|
|
|
|
// Sync UI when the session changes externally (URL back/forward)
|
|
useEffect(() => {
|
|
if (!isRestored) return
|
|
if (!sessionManager.isAvailable) return
|
|
const newSessionId = sessionManager.currentSessionId
|
|
const newSession = sessionManager.currentSession
|
|
if (newSessionId === lastSyncedSessionIdRef.current) return
|
|
const isOwnNewSession =
|
|
!!newSession && savedMessagesRef.current.has(newSession.messages)
|
|
lastSyncedSessionIdRef.current = newSessionId
|
|
if (isOwnNewSession) {
|
|
setPendingTitle(null)
|
|
return
|
|
}
|
|
justLoadedSessionRef.current = true
|
|
syncUIWithSession(newSession)
|
|
}, [
|
|
isRestored,
|
|
sessionManager.isAvailable,
|
|
sessionManager.currentSessionId,
|
|
sessionManager.currentSession,
|
|
syncUIWithSession,
|
|
])
|
|
|
|
const {
|
|
isAvailable: sessionIsAvailable,
|
|
currentSessionId,
|
|
saveCurrentSession,
|
|
getChatGeneration,
|
|
getSaveTicket,
|
|
} = sessionManager
|
|
// Ref avoids re-running the save effect after every save
|
|
const saveCurrentSessionRef = useRef(saveCurrentSession)
|
|
saveCurrentSessionRef.current = saveCurrentSession
|
|
|
|
const versions = useVersionsStore((s) => s.versions)
|
|
|
|
// Debounced auto-save, only when not streaming
|
|
useEffect(() => {
|
|
if (!hasRestoredRef.current) return
|
|
if (!sessionIsAvailable) return
|
|
if (status === "streaming" || status === "submitted") return
|
|
if (justLoadedSessionRef.current) {
|
|
justLoadedSessionRef.current = false
|
|
return
|
|
}
|
|
if (saveDebounceRef.current) clearTimeout(saveDebounceRef.current)
|
|
|
|
// Capture the chat on screen at schedule time; the save is dropped
|
|
// if another chat is on screen by the time it runs
|
|
const scheduledForChat = getChatGeneration()
|
|
const hasDiagramNow = isRealDiagram(chartXMLRef.current)
|
|
const isNoDiagramSession =
|
|
justLoadedSessionIdRef.current === currentSessionId
|
|
|
|
saveDebounceRef.current = setTimeout(async () => {
|
|
try {
|
|
if (messagesRef.current.length > 0 || hasDiagramNow) {
|
|
// Taken before the data is read, for the chat it was
|
|
// scheduled for
|
|
const ticket = {
|
|
...getSaveTicket(),
|
|
generation: scheduledForChat,
|
|
}
|
|
const sessionData = await buildSessionData({
|
|
withThumbnail: hasDiagramNow && !isNoDiagramSession,
|
|
})
|
|
await saveCurrentSessionRef.current(sessionData, ticket)
|
|
}
|
|
} catch (error) {
|
|
console.error("Failed to save session:", error)
|
|
}
|
|
}, SAVE_DEBOUNCE_MS)
|
|
|
|
return () => {
|
|
if (saveDebounceRef.current) clearTimeout(saveDebounceRef.current)
|
|
}
|
|
}, [
|
|
chartXML,
|
|
messages,
|
|
versions,
|
|
status,
|
|
sessionIsAvailable,
|
|
currentSessionId,
|
|
getChatGeneration,
|
|
getSaveTicket,
|
|
buildSessionData,
|
|
])
|
|
|
|
// Put the session in the URL once it exists. Not while the chat is left:
|
|
// the save just gave it a session, and this URL would win over the one
|
|
// of what comes next (another chat, a new one, another language).
|
|
useEffect(() => {
|
|
if (leavingRef.current) return
|
|
if (sessionManager.currentSessionId && !urlSessionId) {
|
|
router.replace(`?session=${sessionManager.currentSessionId}`, {
|
|
scroll: false,
|
|
})
|
|
}
|
|
}, [sessionManager.currentSessionId, urlSessionId, router])
|
|
|
|
// Save when the page is hidden (more reliable than beforeunload for IndexedDB)
|
|
useEffect(() => {
|
|
if (!sessionManager.isAvailable) return
|
|
const handleVisibilityChange = async () => {
|
|
if (
|
|
document.visibilityState === "hidden" &&
|
|
(messagesRef.current.length > 0 ||
|
|
isRealDiagram(chartXMLRef.current))
|
|
) {
|
|
try {
|
|
const ticket = sessionManager.getSaveTicket()
|
|
const sessionData = await buildSessionData({
|
|
withThumbnail: false,
|
|
})
|
|
await sessionManager.saveCurrentSession(sessionData, ticket)
|
|
} catch (error) {
|
|
console.error(
|
|
"Failed to save session on visibility change:",
|
|
error,
|
|
)
|
|
}
|
|
}
|
|
}
|
|
document.addEventListener("visibilitychange", handleVisibilityChange)
|
|
return () =>
|
|
document.removeEventListener(
|
|
"visibilitychange",
|
|
handleVisibilityChange,
|
|
)
|
|
}, [sessionManager, buildSessionData])
|
|
|
|
// ---------------------------------------------------------------------
|
|
// Sending
|
|
// ---------------------------------------------------------------------
|
|
|
|
const isExtractingAttachments =
|
|
files.some((f) => pdfData.get(f)?.isExtracting) ||
|
|
Array.from(urlData.values()).some((d) => d.isExtracting)
|
|
|
|
// Append PDF, text file and URL content to the user's text; images
|
|
// become file parts when imageParts is given
|
|
const processFilesAndAppendContent = async (
|
|
baseText: string,
|
|
imageParts?: any[],
|
|
): Promise<string> => {
|
|
let userText = baseText
|
|
for (const file of files) {
|
|
if (isPdfFile(file)) {
|
|
const extracted = pdfData.get(file)
|
|
if (extracted?.text) {
|
|
userText += `\n\n[PDF: ${file.name}]\n${extracted.text}`
|
|
}
|
|
} else if (isTextFile(file)) {
|
|
const extracted = pdfData.get(file)
|
|
if (extracted?.text) {
|
|
userText += `\n\n[File: ${file.name}]\n${extracted.text}`
|
|
}
|
|
} else if (imageParts) {
|
|
const reader = new FileReader()
|
|
const dataUrl = await new Promise<string>((resolve) => {
|
|
reader.onload = () => resolve(reader.result as string)
|
|
reader.readAsDataURL(file)
|
|
})
|
|
imageParts.push({
|
|
type: "file",
|
|
url: dataUrl,
|
|
mediaType: file.type,
|
|
})
|
|
}
|
|
}
|
|
for (const [url, data] of urlData) {
|
|
if (data.content) {
|
|
userText += `\n\n[URL: ${url}]\nTitle: ${data.title}\n\n${data.content}`
|
|
}
|
|
}
|
|
return userText
|
|
}
|
|
|
|
const getPreviousXml = (beforeIndex: number): string => {
|
|
const snapshotKeys = Array.from(xmlSnapshotsRef.current.keys())
|
|
.filter((k) => k < beforeIndex)
|
|
.sort((a, b) => b - a)
|
|
return snapshotKeys.length > 0
|
|
? getFirstPageXml(
|
|
xmlSnapshotsRef.current.get(snapshotKeys[0]) || "",
|
|
)
|
|
: ""
|
|
}
|
|
|
|
const sendChatMessage = (
|
|
parts: any,
|
|
xml: string,
|
|
previousXml: string,
|
|
turnIndex: number,
|
|
selectedCells: MessageSelection,
|
|
) => {
|
|
autoRetryCountRef.current = 0
|
|
continuationRetryCountRef.current = 0
|
|
validationRetryCountRef.current = 0
|
|
partialXmlRef.current = ""
|
|
continuationOriginalRef.current = null
|
|
stoppedRef.current = false
|
|
currentTurnRef.current = turnIndex
|
|
// Busy from now on, before the next render says so
|
|
busyRef.current = true
|
|
// Its Undo would take this answer's diagram away
|
|
toast.dismiss("version-restored")
|
|
startTurn()
|
|
|
|
const config = getSelectedAIConfig()
|
|
const { customSystemMessage, minimalStyle, maxOutputTokens } =
|
|
useSettingsStore.getState()
|
|
|
|
const selected = selectedCells.map(({ id, label }) => ({ id, label }))
|
|
const options = {
|
|
body: {
|
|
xml,
|
|
previousXml,
|
|
sessionId: langfuseSessionId,
|
|
customSystemMessage,
|
|
...(selected.length > 0 && { selectedCells: selected }),
|
|
},
|
|
headers: buildChatHeaders(config, {
|
|
minimalStyle,
|
|
maxOutputTokens,
|
|
}),
|
|
}
|
|
// The selection stays with the message, for regenerate and retry
|
|
sendMessage(
|
|
{
|
|
parts,
|
|
...(selected.length > 0 && {
|
|
metadata: { selectedCells: selected },
|
|
}),
|
|
},
|
|
options,
|
|
)
|
|
}
|
|
|
|
// Export the current diagram, snapshot it for this message, and send
|
|
const sendWithCurrentDiagram = async (
|
|
parts: any[],
|
|
selectedCells: MessageSelection = [],
|
|
) => {
|
|
const chartXml = formatXML(await onFetchChart())
|
|
const turnIndex = messagesRef.current.length
|
|
const previousXml = getPreviousXml(turnIndex)
|
|
// Snapshot the full multi-page document (kept fresh by autosave) so
|
|
// regenerate/edit can restore every page; the model gets page 1 only
|
|
xmlSnapshotsRef.current.set(turnIndex, chartXMLRef.current || chartXml)
|
|
sendChatMessage(parts, chartXml, previousXml, turnIndex, selectedCells)
|
|
}
|
|
|
|
const clearComposer = () => {
|
|
setInput("")
|
|
setFiles([])
|
|
setUrlData(new Map())
|
|
}
|
|
|
|
const submitInput = async (overrideText?: string) => {
|
|
const text = overrideText ?? input
|
|
if (!text.trim() || isBusy || isExtractingAttachments) return
|
|
|
|
// Cached example: only with no messages and an empty canvas (same
|
|
// rule as the server)
|
|
if (
|
|
messagesRef.current.length === 0 &&
|
|
!hasCells(chartXMLRef.current || "")
|
|
) {
|
|
// The file name keeps a user's own file from matching an example
|
|
const cached = findCachedResponse(
|
|
text.trim(),
|
|
files.length > 0,
|
|
files.length === 1 ? files[0].name : undefined,
|
|
)
|
|
if (cached) {
|
|
const toolCallId = `cached-${Date.now()}`
|
|
const userText = await processFilesAndAppendContent(text)
|
|
currentTurnRef.current = 0
|
|
startTurn()
|
|
setMessages([
|
|
{
|
|
id: `user-${Date.now()}`,
|
|
role: "user" as const,
|
|
parts: [{ type: "text" as const, text: userText }],
|
|
},
|
|
{
|
|
id: `assistant-${Date.now()}`,
|
|
role: "assistant" as const,
|
|
parts: [
|
|
{
|
|
type: "tool-display_diagram" as const,
|
|
toolCallId,
|
|
state: "output-available" as const,
|
|
input: { xml: cached.xml },
|
|
output: "Successfully displayed the diagram.",
|
|
},
|
|
],
|
|
},
|
|
] as any)
|
|
// Snapshot the canvas before the example so editing this message works
|
|
xmlSnapshotsRef.current.set(
|
|
0,
|
|
chartXMLRef.current || BLANK_MXFILE,
|
|
)
|
|
// Load its diagram here: these messages never reach the
|
|
// tool handler
|
|
const prepared = prepareNewDiagram(cached.xml, {
|
|
pageId: "page-1",
|
|
pageName: "Page-1",
|
|
})
|
|
if (prepared.ok) {
|
|
onDisplayChart(prepared.xml, true, "commit", {
|
|
toolCallId,
|
|
})
|
|
}
|
|
clearComposer()
|
|
return
|
|
}
|
|
}
|
|
|
|
const parts: any[] = []
|
|
// The backend only reads the first text part, so combine them
|
|
const userText = await processFilesAndAppendContent(text, parts)
|
|
parts.unshift({ type: "text", text: userText })
|
|
// Clear right away, so text typed while the diagram exports is kept
|
|
const sentFiles = files
|
|
const sentUrls = urlData
|
|
clearComposer()
|
|
setDismissedSelectionKey(selectionKey)
|
|
try {
|
|
await sendWithCurrentDiagram(parts, chatSelection)
|
|
} catch (error) {
|
|
console.error("Error fetching chart data:", error)
|
|
toast.error(dict.errors.failedToExport)
|
|
// Nothing was sent: give the message back, unless something new
|
|
// was typed or attached meanwhile
|
|
if (!inputRef.current && filesRef.current.length === 0) {
|
|
setInput(text)
|
|
setFiles(sentFiles)
|
|
setUrlData(sentUrls)
|
|
}
|
|
}
|
|
}
|
|
|
|
const submit = async (overrideText?: string) => {
|
|
// While a send is prepared (attachments read, diagram exported) the
|
|
// status is still "ready": a second Enter or click would send the
|
|
// message again
|
|
if (preparingSendRef.current || leavingRef.current) return
|
|
preparingSendRef.current = true
|
|
try {
|
|
await submitInput(overrideText)
|
|
} finally {
|
|
preparingSendRef.current = false
|
|
}
|
|
}
|
|
|
|
// Templates fill the input and send it right away (attachments included)
|
|
const sendTemplate = (prompt: string) => {
|
|
setInput(prompt)
|
|
submit(prompt)
|
|
}
|
|
|
|
const improveWithSuggestions = async (feedback: string) => {
|
|
if (!canStartTurn()) return
|
|
// Like submit: the diagram export comes first
|
|
preparingSendRef.current = true
|
|
try {
|
|
await sendWithCurrentDiagram([{ type: "text", text: feedback }])
|
|
} catch (error) {
|
|
console.error("Error fetching chart data:", error)
|
|
toast.error(dict.errors.failedToExport)
|
|
} finally {
|
|
preparingSendRef.current = false
|
|
}
|
|
}
|
|
|
|
const restoreDiagramFromSnapshot = (savedXml: string) => {
|
|
onDisplayChart(savedXml, true)
|
|
chartXMLRef.current = savedXml
|
|
}
|
|
|
|
const cleanupAfterTurn = (messageIndex: number, keepVersions = false) => {
|
|
for (const key of xmlSnapshotsRef.current.keys()) {
|
|
if (key > messageIndex) xmlSnapshotsRef.current.delete(key)
|
|
}
|
|
// The canvas goes back to the snapshot: no kept version is on it
|
|
if (keepVersions) useVersionsStore.getState().clearCanvasFlags()
|
|
else useVersionsStore.getState().removeFromTurn(messageIndex)
|
|
}
|
|
|
|
// Send the user message at index again with these parts, from the
|
|
// diagram it was sent with. keepVersions: a retry keeps what the failed
|
|
// turn drew (in the strip).
|
|
const rerunFrom = (index: number, parts: any[], keepVersions = false) => {
|
|
const savedXml = xmlSnapshotsRef.current.get(index)
|
|
if (!savedXml) {
|
|
console.error("No saved XML snapshot for message index:", index)
|
|
return
|
|
}
|
|
const previousXml = getPreviousXml(index)
|
|
restoreDiagramFromSnapshot(savedXml)
|
|
cleanupAfterTurn(index, keepVersions)
|
|
|
|
// Remove the user message onwards (sendMessage re-adds it). For a
|
|
// moment there may be no message: that is no empty chat.
|
|
flushSync(() => {
|
|
setIsResending(true)
|
|
setMessages(messages.slice(0, index))
|
|
})
|
|
sendChatMessage(
|
|
parts,
|
|
getFirstPageXml(savedXml),
|
|
previousXml,
|
|
index,
|
|
selectionOf(messages[index]),
|
|
)
|
|
}
|
|
|
|
const regenerate = async (messageIndex: number) => {
|
|
if (!canStartTurn()) return
|
|
let userMessageIndex = messageIndex - 1
|
|
while (
|
|
userMessageIndex >= 0 &&
|
|
messages[userMessageIndex].role !== "user"
|
|
) {
|
|
userMessageIndex--
|
|
}
|
|
if (userMessageIndex < 0) return
|
|
const userParts = messages[userMessageIndex].parts
|
|
if (!userParts?.some((p: any) => p.type === "text")) return
|
|
rerunFrom(userMessageIndex, userParts)
|
|
}
|
|
|
|
const editMessage = async (
|
|
messageIndex: number,
|
|
newText: string,
|
|
keepVersions = false,
|
|
) => {
|
|
if (!canStartTurn()) return
|
|
const message = messages[messageIndex]
|
|
if (message?.role !== "user") return
|
|
// The edit box only shows the typed text; keep appended file content
|
|
const newParts = message.parts?.map((part: any) => {
|
|
if (part.type === "text") {
|
|
const appended =
|
|
part.text.match(APPENDED_FILE_SECTIONS_PATTERN)?.[0] ?? ""
|
|
return { ...part, text: newText + appended }
|
|
}
|
|
return part
|
|
}) || [{ type: "text", text: newText }]
|
|
rerunFrom(messageIndex, newParts, keepVersions)
|
|
}
|
|
|
|
/** Send the last user message again (after an error) */
|
|
const retryLastMessage = () => {
|
|
const index = messages.map((m) => m.role).lastIndexOf("user")
|
|
if (index < 0) return
|
|
editMessage(index, getUserOriginalText(messages[index]), true)
|
|
}
|
|
|
|
const handleStop = useCallback(() => {
|
|
stoppedRef.current = true
|
|
stopCountRef.current++
|
|
// A running screenshot check holds up the chat (the SDK waits for
|
|
// the tool handler): end it, so the call gets its result now
|
|
cancelValidation()
|
|
const lastMessage = messages[messages.length - 1]
|
|
// Calls the tool handler already took can still show as streaming:
|
|
// the messages update at most every 150 ms (useChat throttle)
|
|
const toolParts = lastMessage?.parts?.filter(
|
|
(part: any) =>
|
|
part.type?.startsWith("tool-") &&
|
|
part.state === "input-streaming" &&
|
|
!processedToolCallsRef.current.has(part.toolCallId),
|
|
)
|
|
for (const part of (toolParts ?? []) as any[]) {
|
|
if (part.toolCallId) {
|
|
addToolOutput({
|
|
tool: part.type.replace("tool-", ""),
|
|
toolCallId: part.toolCallId,
|
|
state: "output-error",
|
|
errorText: "Stopped by user",
|
|
})
|
|
}
|
|
}
|
|
stop()
|
|
abandonContinuation()
|
|
}, [messages, addToolOutput, stop, cancelValidation])
|
|
|
|
// ---------------------------------------------------------------------
|
|
// Session actions
|
|
// ---------------------------------------------------------------------
|
|
|
|
// The current chat could not be saved (storage full). The list where
|
|
// old chats can be deleted shows only in an empty chat, so let the user
|
|
// go on without saving. It replaces the plain message, and has its own
|
|
// id so a later failed auto-save does not take its button away.
|
|
const offerToContinueUnsaved = useCallback(
|
|
(proceed: () => void) => {
|
|
toast.dismiss("session-save-failed")
|
|
toast.error(dict.errors.sessionSaveFailedLeave, {
|
|
id: "session-save-leave",
|
|
duration: 15000,
|
|
action: {
|
|
label: dict.errors.continueWithoutSaving,
|
|
onClick: proceed,
|
|
},
|
|
})
|
|
},
|
|
[dict],
|
|
)
|
|
|
|
// A new turn makes the offer stale: going on would clear the chat while
|
|
// the answer streams in
|
|
useEffect(() => {
|
|
if (isBusy) toast.dismiss("session-save-leave")
|
|
}, [isBusy])
|
|
|
|
// Save the chat on screen before leaving it (also a diagram drawn
|
|
// without messages). If that failed (storage full), stay on it unless
|
|
// the user goes on without saving: then proceed runs later. Resolves to
|
|
// whether the caller can leave now.
|
|
const saveBeforeLeaving = async (proceed: () => void) => {
|
|
if (!sessionManager.isAvailable) return true
|
|
if (
|
|
messagesRef.current.length === 0 &&
|
|
!isRealDiagram(chartXMLRef.current)
|
|
) {
|
|
return true
|
|
}
|
|
// Of the chat on screen now, also if another one comes on screen
|
|
// while the thumbnail is taken
|
|
const ticket = sessionManager.getSaveTicket()
|
|
const sessionData = await buildSessionData({ withThumbnail: true })
|
|
if (!(await sessionManager.saveCurrentSession(sessionData, ticket))) {
|
|
offerToContinueUnsaved(proceed)
|
|
return false
|
|
}
|
|
return true
|
|
}
|
|
|
|
// Leaving the chat on screen: no answer can start until it is done
|
|
const whileLeaving = async (fn: () => unknown) => {
|
|
leavingRef.current = true
|
|
// The composer shows it: what is typed now could not be sent
|
|
setIsLeaving(true)
|
|
try {
|
|
await fn()
|
|
} finally {
|
|
leavingRef.current = false
|
|
setIsLeaving(false)
|
|
}
|
|
}
|
|
|
|
const selectSession = async (id: string) => {
|
|
if (!sessionManager.isAvailable || !canStartTurn()) return
|
|
const open = async () => {
|
|
if (isTurnActive()) return
|
|
const sessionData = await sessionManager.switchSession(id)
|
|
if (!sessionData) return
|
|
const hasRealDiagram = isRealDiagram(sessionData.diagramXml)
|
|
justLoadedSessionRef.current = true
|
|
// Use the new session's thumbnail, never the previous one's
|
|
latestSvgRef.current = sessionData.thumbnailDataUrl || ""
|
|
justLoadedSessionIdRef.current = hasRealDiagram ? null : id
|
|
// Shown here: the effect that follows session changes must not
|
|
// show it again (that would save it, first in the list)
|
|
lastSyncedSessionIdRef.current = id
|
|
setValidationStates({})
|
|
syncUIWithSession(sessionData)
|
|
router.replace(`?session=${id}`, { scroll: false })
|
|
}
|
|
// Going on without saving (storage full) is a leave of its own
|
|
const proceed = () => whileLeaving(open)
|
|
await whileLeaving(async () => {
|
|
if (await saveBeforeLeaving(proceed)) await open()
|
|
})
|
|
}
|
|
|
|
const deleteSession = useCallback(
|
|
async (id: string) => {
|
|
if (!sessionManager.isAvailable) return
|
|
const isCurrent = id === sessionManager.currentSessionId
|
|
if (isCurrent && !canStartTurn()) return
|
|
const remove = async () => {
|
|
const result = await sessionManager.deleteSession(id)
|
|
if (result.wasCurrentSession) {
|
|
syncUIWithSession(null)
|
|
router.replace(pathname, { scroll: false })
|
|
}
|
|
}
|
|
if (isCurrent) await whileLeaving(remove)
|
|
else await remove()
|
|
},
|
|
[sessionManager, syncUIWithSession, router, pathname, isBusy],
|
|
)
|
|
|
|
// Clear everything for a new conversation (the old one is saved)
|
|
const startNewChat = () => {
|
|
// Clear session state BEFORE the URL, so the URL effect does not
|
|
// bring the old session back
|
|
sessionManager.clearCurrentSession()
|
|
toast.dismiss("version-restored")
|
|
toast.dismiss("session-save-leave")
|
|
// A new conversation brings the start screen back
|
|
useUiStore.getState().setHeroDismissed(false)
|
|
setMessages([])
|
|
setInput("")
|
|
clearDiagram()
|
|
useVersionsStore.getState().clear()
|
|
setValidationStates({})
|
|
handleFileChange([])
|
|
setUrlData(new Map())
|
|
setLangfuseSessionId(newLangfuseSessionId())
|
|
xmlSnapshotsRef.current.clear()
|
|
setPendingTitle(null)
|
|
router.replace(pathname, { scroll: false })
|
|
}
|
|
|
|
const newChat = async () => {
|
|
if (!canStartTurn()) return
|
|
const start = () => {
|
|
if (isTurnActive()) return
|
|
startNewChat()
|
|
toast.success(dict.dialogs.clearSuccess)
|
|
focusComposer()
|
|
}
|
|
await whileLeaving(async () => {
|
|
if (!(await saveBeforeLeaving(start))) return
|
|
// The list shows the chat just saved
|
|
if (sessionManager.isAvailable) {
|
|
await sessionManager.refreshSessions()
|
|
}
|
|
start()
|
|
})
|
|
}
|
|
|
|
const renameSession = useCallback(
|
|
async (title: string) => {
|
|
const id = sessionManager.currentSessionId
|
|
if (!id) {
|
|
// Not saved yet: use the title when the session is created
|
|
setPendingTitle(title.trim().slice(0, 100) || null)
|
|
return
|
|
}
|
|
await sessionManager.renameSession(id, title)
|
|
},
|
|
[sessionManager],
|
|
)
|
|
|
|
const leavePage = async (go: (sessionId: string | null) => void) => {
|
|
if (!canStartTurn()) return
|
|
const proceed = () => go(sessionManager.getCurrentSessionId())
|
|
await whileLeaving(async () => {
|
|
if (await saveBeforeLeaving(proceed)) proceed()
|
|
})
|
|
}
|
|
|
|
const openDiagramFile = async (file: File) => {
|
|
if (!canStartTurn()) return
|
|
await whileLeaving(() => openFile(file))
|
|
}
|
|
|
|
const openFile = async (file: File) => {
|
|
const xml = await readDiagramFile(file).catch(() => null)
|
|
if (!xml) {
|
|
toast.error(
|
|
formatMessage(dict.errors.notADiagramFile, {
|
|
name: file.name,
|
|
}),
|
|
)
|
|
return
|
|
}
|
|
// A loaded file starts a fresh conversation named after it
|
|
const open = () => {
|
|
if (isTurnActive()) return
|
|
startNewChat()
|
|
// This empty chat is ours: the effect that follows session
|
|
// changes must not clear the file it is about to show
|
|
lastSyncedSessionIdRef.current = null
|
|
setPendingTitle(
|
|
file.name.replace(/\.(drawio\.svg|drawio|xml|svg)$/i, ""),
|
|
)
|
|
onDisplayChart(xml, true)
|
|
// The file is on screen, also one without shapes yet
|
|
useUiStore.getState().setHeroDismissed(true)
|
|
}
|
|
// Going on without saving (storage full) is a leave of its own
|
|
const proceed = () => whileLeaving(open)
|
|
if (await saveBeforeLeaving(proceed)) open()
|
|
}
|
|
|
|
const engine: ChatEngine = {
|
|
messages,
|
|
status,
|
|
error,
|
|
isBusy,
|
|
isLeaving,
|
|
isResending,
|
|
setMessages: setMessages as (messages: UIMessage[]) => void,
|
|
stop: handleStop,
|
|
input,
|
|
setInput,
|
|
files,
|
|
pdfData,
|
|
setFiles: handleFileChange,
|
|
urlData,
|
|
setUrlData,
|
|
isExtractingAttachments,
|
|
submit,
|
|
sendTemplate,
|
|
improveWithSuggestions,
|
|
regenerate,
|
|
retryLastMessage,
|
|
editMessage,
|
|
chatSelection,
|
|
dismissSelection,
|
|
sessions: sessionManager.sessions,
|
|
currentSessionId: sessionManager.currentSessionId,
|
|
currentTitle: sessionManager.currentSession?.title ?? pendingTitle,
|
|
isRestored,
|
|
newChat,
|
|
selectSession,
|
|
deleteSession,
|
|
renameSession,
|
|
openDiagramFile,
|
|
leavePage,
|
|
restoreVersion,
|
|
undoVersion,
|
|
langfuseSessionId,
|
|
loadedMessageIdsRef,
|
|
validationStates,
|
|
modelConfig,
|
|
}
|
|
|
|
return (
|
|
<ChatEngineContext.Provider value={engine}>
|
|
{children}
|
|
</ChatEngineContext.Provider>
|
|
)
|
|
}
|