Composable, reusable chat interface with Markdown support + VSCode theme syntax highlighting

webview-ui/package.json

@@ -36,6 +36,7 @@
 		"lucide-react": "^0.475.0",
 		"react": "^18.3.1",
 		"react-dom": "^18.3.1",
+		"react-markdown": "^9.0.3",
 		"react-remark": "^2.1.0",
 		"react-textarea-autosize": "^8.5.3",
 		"react-use": "^17.5.1",
@@ -75,6 +76,7 @@
 		"jest": "^27.5.1",
 		"jest-environment-jsdom": "^27.5.1",
 		"jest-simple-dot-reporter": "^1.0.5",
+		"shiki": "^2.3.2",
 		"storybook": "^8.5.6",
 		"storybook-dark-mode": "^4.0.2",
 		"ts-jest": "^27.1.5",

webview-ui/src/components/ui/chat/Chat.tsx

@@ -0,0 +1,29 @@
+import { HTMLAttributes } from "react"
+
+import { cn } from "@/lib/utils"
+
+import { ChatHandler } from "./types"
+import { ChatProvider } from "./ChatProvider"
+import { ChatMessages } from "./ChatMessages"
+import { ChatInput } from "./ChatInput"
+
+type ChatProps = HTMLAttributes<HTMLDivElement> & {
+	assistantName: string
+	handler: ChatHandler
+}
+
+export const Chat = ({ assistantName, handler, ...props }: ChatProps) => (
+	<ChatProvider value={{ assistantName, ...handler }}>
+		<InnerChat {...props} />
+	</ChatProvider>
+)
+
+type InnerChatProps = HTMLAttributes<HTMLDivElement>
+
+const InnerChat = ({ className, children, ...props }: InnerChatProps) => (
+	<div className={cn("relative flex flex-col flex-1 min-h-0", className)} {...props}>
+		<ChatMessages />
+		{children}
+		<ChatInput />
+	</div>
+)

webview-ui/src/components/ui/chat/ChatInput.tsx

@@ -0,0 +1,100 @@
+import { PaperPlaneIcon, StopIcon } from "@radix-ui/react-icons"
+
+import { Button, AutosizeTextarea } from "@/components/ui"
+
+import { ChatInputProvider } from "./ChatInputProvider"
+import { useChatUI } from "./useChatUI"
+import { useChatInput } from "./useChatInput"
+
+export function ChatInput() {
+	const { input, setInput, append, isLoading } = useChatUI()
+	const isDisabled = isLoading || !input.trim()
+
+	const submit = async () => {
+		if (input.trim() === "") {
+			return
+		}
+
+		setInput("")
+		await append({ role: "user", content: input })
+	}
+
+	const handleSubmit = async (e: React.FormEvent<HTMLFormElement>) => {
+		e.preventDefault()
+		await submit()
+	}
+
+	const handleKeyDown = async (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
+		if (isDisabled) {
+			return
+		}
+
+		if (e.key === "Enter" && !e.shiftKey) {
+			e.preventDefault()
+			await submit()
+		}
+	}
+
+	return (
+		<ChatInputProvider value={{ isDisabled, handleKeyDown, handleSubmit }}>
+			<div className="border-t border-vscode-editor-background p-3">
+				<ChatInputForm />
+			</div>
+		</ChatInputProvider>
+	)
+}
+
+function ChatInputForm() {
+	const { handleSubmit } = useChatInput()
+
+	return (
+		<form onSubmit={handleSubmit} className="relative">
+			<ChatInputField />
+			<ChatInputSubmit />
+		</form>
+	)
+}
+
+interface ChatInputFieldProps {
+	placeholder?: string
+}
+
+function ChatInputField({ placeholder = "Chat" }: ChatInputFieldProps) {
+	const { input, setInput } = useChatUI()
+	const { handleKeyDown } = useChatInput()
+
+	return (
+		<AutosizeTextarea
+			name="input"
+			placeholder={placeholder}
+			minHeight={75}
+			maxHeight={200}
+			value={input}
+			onChange={({ target: { value } }) => setInput(value)}
+			onKeyDown={handleKeyDown}
+			className="resize-none px-3 pt-3 pb-[50px]"
+		/>
+	)
+}
+
+function ChatInputSubmit() {
+	const { isLoading, stop } = useChatUI()
+	const { isDisabled } = useChatInput()
+	const isStoppable = isLoading && !!stop
+
+	return (
+		<div className="absolute bottom-[1px] left-[1px] right-[1px] h-[40px] bg-input border-t border-vscode-editor-background rounded-b-md p-1">
+			<div className="flex flex-row-reverse items-center gap-2">
+				{isStoppable ? (
+					<Button type="button" variant="ghost" size="sm" onClick={stop}>
+						<StopIcon className="text-destructive" />
+					</Button>
+				) : (
+					<Button type="submit" variant="ghost" size="icon" disabled={isDisabled}>
+						<PaperPlaneIcon />
+					</Button>
+				)}
+			</div>
+		</div>
+	)
+}

webview-ui/src/components/ui/chat/ChatInputProvider.ts

@@ -0,0 +1,11 @@
+import { createContext } from "react"
+
+interface ChatInputContext {
+	isDisabled: boolean
+	handleKeyDown: (e: React.KeyboardEvent<HTMLTextAreaElement>) => void
+	handleSubmit: (e: React.FormEvent<HTMLFormElement>) => void
+}
+
+export const chatInputContext = createContext<ChatInputContext | null>(null)
+
+export const ChatInputProvider = chatInputContext.Provider

webview-ui/src/components/ui/chat/ChatMessage.tsx

@@ -0,0 +1,105 @@
+import { useMemo } from "react"
+import { CopyIcon, CheckIcon } from "@radix-ui/react-icons"
+import { BrainCircuit, CircleUserRound } from "lucide-react"
+
+import { cn } from "@/lib/utils"
+import { useClipboard } from "@/components/ui/hooks"
+import { Badge } from "@/components/ui"
+import { Markdown } from "@/components/ui/markdown"
+
+import { BadgeData, ChatHandler, Message, MessageAnnotationType } from "./types"
+import { ChatMessageProvider } from "./ChatMessageProvider"
+import { useChatUI } from "./useChatUI"
+import { useChatMessage } from "./useChatMessage"
+
+interface ChatMessageProps {
+	message: Message
+	isLast: boolean
+	isHeaderVisible: boolean
+	isLoading?: boolean
+	append?: ChatHandler["append"]
+}
+
+export function ChatMessage({ message, isLast, isHeaderVisible, isLoading, append }: ChatMessageProps) {
+	const badges = useMemo(
+		() =>
+			message.annotations
+				?.filter(({ type }) => type === MessageAnnotationType.BADGE)
+				.map(({ data }) => data as BadgeData),
+		[message.annotations],
+	)
+
+	return (
+		<ChatMessageProvider value={{ message, isLast }}>
+			<div
+				className={cn("relative group flex flex-col text-secondary-foreground", {
+					"bg-vscode-input-background/50": message.role === "user",
+				})}>
+				{isHeaderVisible && <ChatMessageHeader badges={badges} />}
+				<ChatMessageContent isHeaderVisible={isHeaderVisible} />
+				<ChatMessageActions />
+			</div>
+		</ChatMessageProvider>
+	)
+}
+
+interface ChatMessageHeaderProps {
+	badges?: BadgeData[]
+}
+
+function ChatMessageHeader({ badges }: ChatMessageHeaderProps) {
+	return (
+		<div className="flex flex-row items-center justify-between border-t border-accent px-3 pt-3 pb-1">
+			<ChatMessageAvatar />
+			{badges?.map(({ label, variant = "outline" }) => (
+				<Badge variant={variant} key={label}>
+					{label}
+				</Badge>
+			))}
+		</div>
+	)
+}
+
+const icons: Record<string, React.ReactNode> = {
+	user: <CircleUserRound className="h-4 w-4" />,
+	assistant: <BrainCircuit className="h-4 w-4" />,
+}
+
+function ChatMessageAvatar() {
+	const { assistantName } = useChatUI()
+	const { message } = useChatMessage()
+
+	return icons[message.role] ? (
+		<div className="flex flex-row items-center gap-1">
+			<div className="opacity-25 select-none">{icons[message.role]}</div>
+			<div className="text-muted">{message.role === "user" ? "You" : assistantName}</div>
+		</div>
+	) : null
+}
+
+interface ChatMessageContentProps {
+	isHeaderVisible: boolean
+}
+
+function ChatMessageContent({ isHeaderVisible }: ChatMessageContentProps) {
+	const { message } = useChatMessage()
+
+	return (
+		<div className={cn("flex flex-col gap-4 flex-1 min-w-0 px-4 pb-6", { "pt-4": isHeaderVisible })}>
+			<Markdown content={message.content} />
+		</div>
+	)
+}
+
+function ChatMessageActions() {
+	const { message } = useChatMessage()
+	const { isCopied, copy } = useClipboard()
+
+	return (
+		<div
+			className="absolute right-2 bottom-2 opacity-0 group-hover:opacity-25 cursor-pointer"
+			onClick={() => copy(message.content)}>
+			{isCopied ? <CheckIcon /> : <CopyIcon />}
+		</div>
+	)
+}

webview-ui/src/components/ui/chat/ChatMessageProvider.ts

@@ -0,0 +1,12 @@
+import { createContext } from "react"
+
+import { Message } from "./types"
+
+export interface ChatMessageContext {
+	message: Message
+	isLast: boolean
+}
+
+export const chatMessageContext = createContext<ChatMessageContext | null>(null)
+
+export const ChatMessageProvider = chatMessageContext.Provider

webview-ui/src/components/ui/chat/ChatMessages.tsx

@@ -0,0 +1,42 @@
+import { useCallback, useEffect, useRef } from "react"
+
+import { useChatUI } from "./useChatUI"
+import { ChatMessage } from "./ChatMessage"
+
+export function ChatMessages() {
+	const { messages, isLoading, append } = useChatUI()
+	const containerRef = useRef<HTMLDivElement>(null)
+	const messageCount = messages.length
+
+	const scrollToBottom = useCallback(() => {
+		if (!containerRef.current) {
+			return
+		}
+
+		requestAnimationFrame(() => {
+			containerRef.current?.scrollTo({
+				top: containerRef.current.scrollHeight,
+				behavior: "smooth",
+			})
+		})
+	}, [])
+
+	useEffect(() => scrollToBottom(), [messageCount, scrollToBottom])
+
+	return (
+		<div ref={containerRef} className="flex flex-col flex-1 min-h-0 overflow-auto relative">
+			{messages.map((message, index) => (
+				<ChatMessage
+					key={index}
+					message={message}
+					isHeaderVisible={
+						!!message.annotations?.length || index === 0 || messages[index - 1].role !== message.role
+					}
+					isLast={index === messageCount - 1}
+					isLoading={isLoading}
+					append={append}
+				/>
+			))}
+		</div>
+	)
+}

webview-ui/src/components/ui/chat/ChatProvider.ts

@@ -0,0 +1,11 @@
+import { createContext } from "react"
+
+import { ChatHandler } from "./types"
+
+type ChatContext = ChatHandler & {
+	assistantName: string
+}
+
+export const chatContext = createContext<ChatContext | null>(null)
+
+export const ChatProvider = chatContext.Provider

webview-ui/src/components/ui/chat/index.ts

@@ -0,0 +1,2 @@
+export * from "./types"
+export * from "./Chat"

webview-ui/src/components/ui/chat/types.ts

@@ -0,0 +1,39 @@
+export interface Message {
+	role: "system" | "user" | "assistant" | "data"
+	content: string
+	annotations?: MessageAnnotation[]
+}
+
+export type ChatHandler = {
+	isLoading: boolean
+	setIsLoading: (isLoading: boolean, message?: string) => void
+
+	loadingMessage?: string
+	setLoadingMessage?: (message: string) => void
+
+	input: string
+	setInput: (input: string) => void
+
+	messages: Message[]
+
+	reload?: (options?: { data?: any }) => void
+	stop?: () => void
+	append: (message: Message, options?: { data?: any }) => Promise<string | null | undefined>
+	reset?: () => void
+}
+
+export enum MessageAnnotationType {
+	BADGE = "badge",
+}
+
+export type BadgeData = {
+	label: string
+	variant?: "default" | "secondary" | "destructive" | "outline"
+}
+
+export type AnnotationData = BadgeData
+
+export type MessageAnnotation = {
+	type: MessageAnnotationType
+	data: AnnotationData
+}

webview-ui/src/components/ui/chat/useChatInput.ts

@@ -0,0 +1,13 @@
+import { useContext } from "react"
+
+import { chatInputContext } from "./ChatInputProvider"
+
+export const useChatInput = () => {
+	const context = useContext(chatInputContext)
+
+	if (!context) {
+		throw new Error("useChatInput must be used within a ChatInputProvider")
+	}
+
+	return context
+}

webview-ui/src/components/ui/chat/useChatMessage.ts

@@ -0,0 +1,13 @@
+import { useContext } from "react"
+
+import { chatMessageContext } from "./ChatMessageProvider"
+
+export const useChatMessage = () => {
+	const context = useContext(chatMessageContext)
+
+	if (!context) {
+		throw new Error("useChatMessage must be used within a ChatMessageProvider")
+	}
+
+	return context
+}

webview-ui/src/components/ui/chat/useChatUI.ts

@@ -0,0 +1,13 @@
+import { useContext } from "react"
+
+import { chatContext } from "./ChatProvider"
+
+export const useChatUI = () => {
+	const context = useContext(chatContext)
+
+	if (!context) {
+		throw new Error("useChatUI must be used within a ChatProvider")
+	}
+
+	return context
+}

webview-ui/src/components/ui/hooks/index.ts

@@ -0,0 +1 @@
+export * from "./useClipboard"

webview-ui/src/components/ui/hooks/useClipboard.ts

@@ -0,0 +1,22 @@
+import { useState } from "react"
+
+export interface UseClipboardProps {
+	timeout?: number
+}
+
+export function useClipboard({ timeout = 2000 }: UseClipboardProps = {}) {
+	const [isCopied, setIsCopied] = useState(false)
+
+	const copy = (value: string) => {
+		if (typeof window === "undefined" || !navigator.clipboard?.writeText || !value) {
+			return
+		}
+
+		navigator.clipboard.writeText(value).then(() => {
+			setIsCopied(true)
+			setTimeout(() => setIsCopied(false), timeout)
+		})
+	}
+
+	return { isCopied, copy }
+}

webview-ui/src/components/ui/markdown/Blockquote.tsx

@@ -0,0 +1,3 @@
+export const Blockquote = ({ children }: { children: React.ReactNode }) => {
+	return <div className="border-l-3 border-accent italic pl-2 py-2 mb-2">{children}</div>
+}

webview-ui/src/components/ui/markdown/CodeBlock.tsx

@@ -0,0 +1,72 @@
+import { FC, memo, useState, useEffect, useCallback } from "react"
+import { codeToHtml } from "shiki"
+import { CopyIcon, CheckIcon } from "@radix-ui/react-icons"
+
+import { cn } from "@/lib/utils"
+import { useClipboard } from "@/components/ui/hooks"
+import { Button } from "@/components/ui"
+
+interface CodeBlockProps extends Omit<React.HTMLAttributes<HTMLDivElement>, "children"> {
+	language: string
+	value: string
+}
+
+export const CodeBlock: FC<CodeBlockProps> = memo(({ language, value, className, ...props }) => {
+	const [highlightedCode, setHighlightedCode] = useState<string>("")
+	const { isCopied, copy } = useClipboard()
+
+	const onCopy = useCallback(() => {
+		if (!isCopied) {
+			copy(value)
+		}
+	}, [isCopied, copy, value])
+
+	useEffect(() => {
+		const highlight = async () => {
+			const theme = "github-dark" // Use VSCode's current theme.
+
+			try {
+				const html = await codeToHtml(value, {
+					lang: language,
+					theme,
+					transformers: [
+						{
+							pre(node) {
+								node.properties.class = cn(className, "overflow-x-auto")
+								return node
+							},
+							code(node) {
+								node.properties.style = "background-color: transparent !important;"
+								return node
+							},
+						},
+					],
+				})
+
+				setHighlightedCode(html)
+			} catch (e) {
+				setHighlightedCode(value)
+			}
+		}
+
+		highlight()
+	}, [language, value, className])
+
+	return (
+		<div className="relative" {...props}>
+			<div dangerouslySetInnerHTML={{ __html: highlightedCode }} />
+			<Button
+				variant="outline"
+				size="icon"
+				className="absolute top-1 right-1 cursor-pointer bg-black/10"
+				onClick={onCopy}>
+				{isCopied ? (
+					<CheckIcon style={{ width: 12, height: 12 }} />
+				) : (
+					<CopyIcon style={{ width: 12, height: 12 }} />
+				)}
+			</Button>
+		</div>
+	)
+})
+CodeBlock.displayName = "CodeBlock"

webview-ui/src/components/ui/markdown/Markdown.tsx

@@ -0,0 +1,107 @@
+import { FC, memo } from "react"
+import ReactMarkdown, { Options } from "react-markdown"
+
+import { Separator } from "@/components/ui"
+
+import { CodeBlock } from "./CodeBlock"
+import { SourceNumberButton } from "./SourceNumberButton"
+import { Blockquote } from "./Blockquote"
+
+const MemoizedReactMarkdown: FC<Options> = memo(
+	ReactMarkdown,
+	(prevProps, nextProps) => prevProps.children === nextProps.children && prevProps.className === nextProps.className,
+)
+
+const preprocessLaTeX = (content: string) => {
+	// Replace block-level LaTeX delimiters \[ \] with $$ $$
+	const blockProcessedContent = content.replace(/\\\[([\s\S]*?)\\\]/g, (_, equation) => `$$${equation}$$`)
+
+	// Replace inline LaTeX delimiters \( \) with $ $
+	return blockProcessedContent.replace(/\\\[([\s\S]*?)\\\]/g, (_, equation) => `$${equation}$`)
+}
+
+export function Markdown({ content }: { content: string }) {
+	const processedContent = preprocessLaTeX(content)
+
+	return (
+		<MemoizedReactMarkdown
+			className="custom-markdown break-words"
+			components={{
+				p({ children }) {
+					return <div className="mb-2 last:mb-0">{children}</div>
+				},
+				hr() {
+					return <Separator />
+				},
+				ol({ children }) {
+					return (
+						<ol className="list-decimal pl-4 [&>li]:mb-1 [&>li:last-child]:mb-0 [&>li>ul]:mt-1 [&>li>ol]:mt-1">
+							{children}
+						</ol>
+					)
+				},
+				ul({ children }) {
+					return (
+						<ul className="list-disc pl-4 [&>li]:mb-1 [&>li:last-child]:mb-0 [&>li>ul]:mt-1 [&>li>ol]:mt-1">
+							{children}
+						</ul>
+					)
+				},
+				blockquote({ children }) {
+					return <Blockquote>{children}</Blockquote>
+				},
+				code({ className, children, ...props }) {
+					if (children && Array.isArray(children) && children.length) {
+						if (children[0] === "▍") {
+							return <span className="mt-1 animate-pulse cursor-default">▍</span>
+						}
+
+						children[0] = (children[0] as string).replace("`▍`", "▍")
+					}
+
+					const match = /language-(\w+)/.exec(className || "")
+
+					const isInline =
+						props.node?.position && props.node.position.start.line === props.node.position.end.line
+
+					return isInline ? (
+						<code className={className} {...props}>
+							{children}
+						</code>
+					) : (
+						<CodeBlock
+							language={(match && match[1]) || ""}
+							value={String(children).replace(/\n$/, "")}
+							className="rounded-xs p-3 mb-2"
+						/>
+					)
+				},
+				a({ href, children }) {
+					// If a text link starts with 'citation:', then render it as
+					// a citation reference.
+					if (
+						Array.isArray(children) &&
+						typeof children[0] === "string" &&
+						children[0].startsWith("citation:")
+					) {
+						const index = Number(children[0].replace("citation:", ""))
+
+						if (!isNaN(index)) {
+							return <SourceNumberButton index={index} />
+						}
+
+						// Citation is not looked up yet, don't render anything.
+						return null
+					}
+
+					return (
+						<a href={href} target="_blank" rel="noopener noreferrer">
+							{children}
+						</a>
+					)
+				},
+			}}>
+			{processedContent}
+		</MemoizedReactMarkdown>
+	)
+}

webview-ui/src/components/ui/markdown/SourceNumberButton.tsx

@@ -0,0 +1,13 @@
+import { cn } from "@/lib/utils"
+
+export function SourceNumberButton({ index, className }: { index: number; className?: string }) {
+	return (
+		<span
+			className={cn(
+				"inline-flex h-5 w-5 items-center justify-center rounded-full bg-gray-100 text-xs",
+				className,
+			)}>
+			{index + 1}
+		</span>
+	)
+}

webview-ui/src/components/ui/markdown/index.ts

@@ -0,0 +1 @@
+export * from "./Markdown"

webview-ui/src/index.css

@@ -358,3 +358,11 @@ vscode-dropdown::part(listbox) {
 input[cmdk-input]:focus {
 	outline: none;
 }
+
+/**
+ * Markdown
+ */
+
+.custom-markdown > pre {
+	background-color: transparent !important;
+}

webview-ui/src/stories/Chat.stories.tsx

@@ -0,0 +1,51 @@
+import { useState } from "react"
+import type { Meta, StoryObj } from "@storybook/react"
+
+import { Chat, ChatHandler, Message } from "@/components/ui/chat"
+
+const meta = {
+	title: "ui/Chat",
+	component: Chat,
+	parameters: { layout: "centered" },
+	tags: ["autodocs"],
+} satisfies Meta<typeof Chat>
+
+export default meta
+
+type Story = StoryObj<typeof meta>
+
+export const Default: Story = {
+	name: "Chat",
+	args: {
+		assistantName: "Assistant",
+		handler: {} as ChatHandler,
+	},
+	render: function StorybookChat() {
+		const handler = useStorybookChat()
+		return (
+			<Chat
+				assistantName="Assistant"
+				handler={handler}
+				className="border w-[460px] h-[640px] bg-vscode-editor-background"
+			/>
+		)
+	},
+}
+
+const useStorybookChat = (): ChatHandler => {
+	const [isLoading, setIsLoading] = useState(false)
+	const [input, setInput] = useState("")
+	const [messages, setMessages] = useState<Message[]>([])
+
+	const append = async (message: Message, options?: { data?: any }) => {
+		const echo: Message = {
+			...message,
+			role: "assistant",
+			content: "Lorem ipsum dolor sit amet, consectetur adipiscing elit.",
+		}
+		setMessages((prev) => [...prev, message, echo])
+		return Promise.resolve(null)
+	}
+
+	return { isLoading, setIsLoading, input, setInput, messages, append }
+}

webview-ui/vite.config.ts

@@ -35,5 +35,6 @@ export default defineConfig({
 	},
 	define: {
 		"process.platform": JSON.stringify(process.platform),
+		"process.env.VSCODE_TEXTMATE_DEBUG": JSON.stringify(process.env.VSCODE_TEXTMATE_DEBUG),
 	},
 })