diff --git a/frontend/app/aipanel/aimessage.tsx b/frontend/app/aipanel/aimessage.tsx index 34c4773f4b..26386cb090 100644 --- a/frontend/app/aipanel/aimessage.tsx +++ b/frontend/app/aipanel/aimessage.tsx @@ -1,11 +1,13 @@ // Copyright 2025, Command Line Inc. // SPDX-License-Identifier: Apache-2.0 +import { WaveStreamdown } from "@/app/element/streamdown"; import { cn } from "@/util/util"; +import { useAtomValue } from "jotai"; import { memo } from "react"; -import { Streamdown } from "streamdown"; import { getFileIcon } from "./ai-utils"; import { WaveUIMessage, WaveUIMessagePart } from "./aitypes"; +import { WaveAIModel } from "./waveai-model"; const AIThinking = memo(() => (
@@ -72,6 +74,8 @@ interface AIMessagePartProps { } const AIMessagePart = memo(({ part, role, isStreaming }: AIMessagePartProps) => { + const model = WaveAIModel.getInstance(); + if (part.type === "text") { const content = part.text ?? ""; @@ -79,23 +83,12 @@ const AIMessagePart = memo(({ part, role, isStreaming }: AIMessagePartProps) => return
{content}
; } else { return ( - - {content} - + className="text-gray-100" + codeBlockMaxWidthAtom={model.codeBlockMaxWidth} + /> ); } } @@ -139,9 +132,7 @@ export const AIMessage = memo(({ message, isStreaming }: AIMessageProps) => {
{showThinkingOnly ? ( diff --git a/frontend/app/aipanel/aipanel.tsx b/frontend/app/aipanel/aipanel.tsx index a5d30d8440..d92c029130 100644 --- a/frontend/app/aipanel/aipanel.tsx +++ b/frontend/app/aipanel/aipanel.tsx @@ -34,6 +34,7 @@ const AIPanelComponentInner = memo(({ className, onClose }: AIPanelProps) => { const [isDragOver, setIsDragOver] = useState(false); const [isLoadingChat, setIsLoadingChat] = useState(true); const model = WaveAIModel.getInstance(); + const containerRef = useRef(null); const errorMessage = jotai.useAtomValue(model.errorMessage); const realMessageRef = useRef(null); const inputRef = useRef(null); @@ -104,10 +105,32 @@ const AIPanelComponentInner = memo(({ className, onClose }: AIPanelProps) => { const messages = await model.loadChat(); setMessages(messages as any); setIsLoadingChat(false); + setTimeout(() => { + model.scrollToBottom(); + }, 100); }; loadMessages(); }, [model, setMessages]); + useEffect(() => { + const updateWidth = () => { + if (containerRef.current) { + globalStore.set(model.containerWidth, containerRef.current.offsetWidth); + } + }; + + updateWidth(); + + const resizeObserver = new ResizeObserver(updateWidth); + if (containerRef.current) { + resizeObserver.observe(containerRef.current); + } + + return () => { + resizeObserver.disconnect(); + }; + }, [model]); + useEffect(() => { model.ensureRateLimitSet(); }, [model]); @@ -279,6 +302,7 @@ const AIPanelComponentInner = memo(({ className, onClose }: AIPanelProps) => { return (
{ return ( @@ -25,6 +26,7 @@ interface AIPanelMessagesProps { } export const AIPanelMessages = memo(({ messages, status, isLoadingChat }: AIPanelMessagesProps) => { + const model = WaveAIModel.getInstance(); const isPanelOpen = useAtomValue(WorkspaceLayoutModel.getInstance().panelVisibleAtom); const messagesEndRef = useRef(null); const messagesContainerRef = useRef(null); @@ -37,6 +39,10 @@ export const AIPanelMessages = memo(({ messages, status, isLoadingChat }: AIPane } }; + useEffect(() => { + model.registerScrollToBottom(scrollToBottom); + }, [model]); + useEffect(() => { scrollToBottom(); }, [messages]); diff --git a/frontend/app/aipanel/waveai-model.tsx b/frontend/app/aipanel/waveai-model.tsx index 36f0ecd9eb..852d0a11ae 100644 --- a/frontend/app/aipanel/waveai-model.tsx +++ b/frontend/app/aipanel/waveai-model.tsx @@ -24,12 +24,15 @@ export interface DroppedFile { export class WaveAIModel { private static instance: WaveAIModel | null = null; private inputRef: React.RefObject | null = null; + private scrollToBottomCallback: (() => void) | null = null; widgetAccessAtom!: jotai.Atom; droppedFiles: jotai.PrimitiveAtom = jotai.atom([]); chatId!: jotai.PrimitiveAtom; errorMessage: jotai.PrimitiveAtom = jotai.atom(null) as jotai.PrimitiveAtom; modelAtom!: jotai.Atom; + containerWidth: jotai.PrimitiveAtom = jotai.atom(0); + codeBlockMaxWidth!: jotai.Atom; private constructor() { const tabId = globalStore.get(atoms.staticTabId); @@ -58,6 +61,11 @@ export class WaveAIModel { const value = get(widgetAccessMetaAtom); return value ?? true; }); + + this.codeBlockMaxWidth = jotai.atom((get) => { + const width = get(this.containerWidth); + return width > 0 ? width - 35 : 0; + }); } static getInstance(): WaveAIModel { @@ -140,6 +148,14 @@ export class WaveAIModel { this.inputRef = ref; } + registerScrollToBottom(callback: () => void) { + this.scrollToBottomCallback = callback; + } + + scrollToBottom() { + this.scrollToBottomCallback?.(); + } + focusInput() { if (!WorkspaceLayoutModel.getInstance().getAIPanelVisible()) { WorkspaceLayoutModel.getInstance().setAIPanelVisible(true); diff --git a/frontend/app/element/streamdown.tsx b/frontend/app/element/streamdown.tsx new file mode 100644 index 0000000000..ba848ee904 --- /dev/null +++ b/frontend/app/element/streamdown.tsx @@ -0,0 +1,317 @@ +// Copyright 2025, Command Line Inc. +// SPDX-License-Identifier: Apache-2.0 + +import { CopyButton } from "@/app/element/copybutton"; +import { IconButton } from "@/app/element/iconbutton"; +import { cn, useAtomValueSafe } from "@/util/util"; +import type { Atom } from "jotai"; +import { useCallback, useEffect, useMemo, useRef, useState } from "react"; +import { bundledLanguages, codeToHtml } from "shiki/bundle/web"; +import { Streamdown } from "streamdown"; +import { throttle } from "throttle-debounce"; + +const ShikiTheme = "github-dark-high-contrast"; + +function extractText(node: React.ReactNode): string { + if (node == null || typeof node === "boolean") return ""; + if (typeof node === "string" || typeof node === "number") return String(node); + if (Array.isArray(node)) return node.map(extractText).join(""); + // @ts-expect-error props exists on ReactElement + if (typeof node === "object" && node.props) return extractText(node.props.children); + return ""; +} + +function CodePlain({ className = "", isCodeBlock, text }: { className?: string; isCodeBlock: boolean; text: string }) { + if (isCodeBlock) { + return {text}; + } + + return ( + + {text} + + ); +} + +function CodeHighlight({ className = "", lang, text }: { className?: string; lang: string; text: string }) { + const [html, setHtml] = useState(""); + const [hasError, setHasError] = useState(false); + const codeRef = useRef(null); + const seqRef = useRef(0); + + const highlightCode = useCallback( + async (textToHighlight: string, language: string, disposedRef: { current: boolean }, seq: number) => { + try { + const full = await codeToHtml(textToHighlight, { lang: language, theme: ShikiTheme }); + const start = full.indexOf("", start); + const end = full.lastIndexOf(""); + const inner = start !== -1 && open !== -1 && end !== -1 ? full.slice(open + 1, end) : ""; + if (!disposedRef.current && seq === seqRef.current) { + setHtml(inner); + setHasError(false); + } + } catch (e) { + if (!disposedRef.current && seq === seqRef.current) { + setHasError(true); + } + console.warn(`Shiki highlight failed for ${language}`, e); + } + }, + [] + ); + + const throttledHighlight = useMemo(() => throttle(300, highlightCode, { noLeading: false }), [highlightCode]); + + useEffect(() => { + const disposedRef = { current: false }; + + if (!text) { + setHtml(""); + return; + } + + seqRef.current++; + const currentSeq = seqRef.current; + throttledHighlight(text, lang, disposedRef, currentSeq); + + return () => { + disposedRef.current = true; + }; + }, [text, lang, throttledHighlight]); + + if (hasError) { + return ( + + {text} + + ); + } + + if (!html && text) { + return ( + + {text} + + ); + } + + return ( + + ); +} + +export function Code({ className = "", children }: { className?: string; children: React.ReactNode }) { + const m = className?.match(/language-([\w+-]+)/i); + const isCodeBlock = !!m; + const lang = m?.[1] || "text"; + const text = extractText(children); + + if (isCodeBlock && lang in bundledLanguages) { + return ; + } + + return ; +} + +type CodeBlockProps = { + children: React.ReactNode; + onClickExecute?: (cmd: string) => void; + codeBlockMaxWidthAtom?: Atom; +}; + +const CodeBlock = ({ children, onClickExecute, codeBlockMaxWidthAtom }: CodeBlockProps) => { + const codeBlockMaxWidth = useAtomValueSafe(codeBlockMaxWidthAtom); + const getLanguage = (children: any): string => { + if (children?.props?.className) { + const match = children.props.className.match(/language-([\w+-]+)/i); + if (match) return match[1]; + } + return "text"; + }; + + const handleCopy = async (e: React.MouseEvent) => { + const textToCopy = extractText(children).replace(/\n$/, ""); + await navigator.clipboard.writeText(textToCopy); + }; + + const handleExecute = (e: React.MouseEvent) => { + const cmd = extractText(children).replace(/\n$/, ""); + if (onClickExecute) { + onClickExecute(cmd); + return; + } + }; + + const language = getLanguage(children); + + return ( +
+
+ {language} +
+ + {onClickExecute && ( + + )} +
+
+ 
{children}
+
+ ); +}; + +function Collapsible({ title, children, defaultOpen = false }) { + const [isOpen, setIsOpen] = useState(defaultOpen); + + return ( +
+ + {isOpen &&
{children}
} +
+ ); +} + +interface WaveStreamdownProps { + text: string; + parseIncompleteMarkdown?: boolean; + className?: string; + onClickExecute?: (cmd: string) => void; + codeBlockMaxWidthAtom?: Atom; +} + +export const WaveStreamdown = ({ + text, + parseIncompleteMarkdown, + className, + onClickExecute, + codeBlockMaxWidthAtom, +}: WaveStreamdownProps) => { + const components = useMemo( + () => ({ + code: Code, + pre: (props: React.HTMLAttributes) => ( + + ), + p: (props: React.HTMLAttributes) =>

, + h1: (props: React.HTMLAttributes) => ( +

+ ), + h2: (props: React.HTMLAttributes) => ( +

+ ), + h3: (props: React.HTMLAttributes) => ( +

+ ), + h4: (props: React.HTMLAttributes) => ( +

+ ), + h5: (props: React.HTMLAttributes) => ( +
+ ), + h6: (props: React.HTMLAttributes) => ( +
+ ), + table: (props: React.HTMLAttributes) => ( + + ), + thead: (props: React.HTMLAttributes) => ( + + ), + tbody: (props: React.HTMLAttributes) => , + tr: (props: React.HTMLAttributes) => ( + + ), + th: (props: React.HTMLAttributes) => ( +
+ ), + td: (props: React.HTMLAttributes) => ( + + ), + ul: (props: React.HTMLAttributes) => ( +