feat(ai-search): enhance AI search panel with improved layout and persistence

- Refactor AI search components into separate panel and trigger elements
- Add localStorage persistence for search input
- Improve button styling and variants
- Optimize markdown rendering performance
- Fix error handling in copy button functionality

Author: ryan-zayne

apps/docs/app/(home)/page.tsx

@@ -1,7 +1,7 @@
-import Link from "next/link";
 import { DocsIcon, GitHubIcon } from "@/components/icons";
 import Ripple from "@/components/landing/ripple";
 import { Button } from "@/components/ui/button";
+import Link from "next/link";
 
 export default function HomePage() {
 	return (
@@ -20,14 +20,14 @@ export default function HomePage() {
 
 				<div className="flex w-full items-center justify-center gap-4 py-4">
 					<Link href="/docs">
-						<Button className="flex gap-2">
+						<Button className="flex gap-2" size="home-default">
 							<DocsIcon />
 							Docs
 						</Button>
 					</Link>
 
 					<Link href="https://github.com/zayne-labs/callapi">
-						<Button className="flex gap-2" theme="secondary">
+						<Button className="flex gap-2" theme="secondary" size="home-default">
 							<GitHubIcon />
 							Github
 						</Button>

apps/docs/app/docs/layout.tsx

@@ -1,4 +1,4 @@
-import { AISearchTrigger } from "@/components/ai/search";
+import { AISearchPanel, AISearchRoot, AISearchTrigger } from "@/components/ai/search";
 import { BgPattern } from "@/components/icons";
 import "fumadocs-twoslash/twoslash.css";
 import { DocsLayout } from "fumadocs-ui/layouts/notebook";
@@ -18,7 +18,10 @@ function Layout(props: LayoutProps<"/docs">) {
 
 			{children}
 
-			<AISearchTrigger />
+			<AISearchRoot>
+				<AISearchPanel />
+				<AISearchTrigger />
+			</AISearchRoot>
 		</DocsLayout>
 	);
 }

apps/docs/components/ai/markdown.tsx

@@ -1,3 +1,4 @@
+import { tw } from "@zayne-labs/toolkit-core";
 import { DynamicCodeBlock } from "fumadocs-ui/components/dynamic-codeblock";
 import defaultMdxComponents from "fumadocs-ui/mdx";
 import type { ElementContent, Root, RootContent } from "hast";
@@ -42,7 +43,7 @@ function rehypeWrapWords() {
 				return {
 					children: [{ type: "text", value: word }],
 					properties: {
-						class: "animate-fd-fade-in",
+						class: tw`animate-fd-fade-in`,
 					},
 					tagName: "span",
 					type: "element",
@@ -110,19 +111,23 @@ function Pre(props: ComponentProps<"pre">) {
 
 const processor = createProcessor();
 
-export function Markdown({ text }: { text: string }) {
+export function Markdown(props: { text: string }) {
+	const { text } = props;
+
 	const deferredText = useDeferredValue(text);
 
 	return (
-		<Suspense fallback={text}>
+		<Suspense fallback={<p className="invisible">{text}</p>}>
 			<Renderer text={deferredText} />
 		</Suspense>
 	);
 }
 
 const cache = new Map<string, Promise<ReactNode>>();
 
-function Renderer({ text }: { text: string }) {
+function Renderer(props: { text: string }) {
+	const { text } = props;
+
 	const result = cache.get(text) ?? processor.process(text);
 	cache.set(text, result);
 

apps/docs/components/ai/page-actions.tsx

@@ -32,9 +32,13 @@ export function LLMCopyButton(props: CopyBtnProps) {
 		const markDownPromise = callApi(markdownUrl, {
 			responseType: "text",
 		}).then((result) => {
-			const content = result.data ?? "";
-			Boolean(content) && cache.set(markdownUrl, content);
-			return content;
+			if (result.error || !result.data) {
+				return "";
+			}
+
+			cache.set(markdownUrl, result.data);
+
+			return result.data;
 		});
 
 		const clipboardItem = new ClipboardItem({ "text/plain": markDownPromise });
@@ -58,7 +62,7 @@ export function LLMCopyButton(props: CopyBtnProps) {
 			{checked ?
 				<Check />
 			:	<Copy />}
-			<p>Copy Markdown</p>
+			Copy Markdown
 		</Button>
 	);
 }
@@ -88,12 +92,12 @@ export function ViewOptions(props: ViewOptionsProps) {
 		<Popover>
 			<PopoverTrigger asChild={true}>
 				<Button theme="secondary" size="sm" className="gap-2">
-					<p>Open</p>
+					Open
 					<ChevronDown className="size-3.5 text-fd-muted-foreground" />
 				</Button>
 			</PopoverTrigger>
 
-			<PopoverContent className="flex flex-col overflow-auto">
+			<PopoverContent className="flex flex-col">
 				{items.map((item) => (
 					<a
 						key={item.href}

apps/docs/components/ai/search.tsx

@@ -10,7 +10,7 @@ import { Presence } from "@zayne-labs/ui-react/common/presence";
 import { DefaultChatTransport } from "ai";
 import Link from "fumadocs-core/link";
 import { Loader2, MessageCircleIcon, RefreshCw, Send, X } from "lucide-react";
-import { useEffect, useMemo, useRef, useState } from "react";
+import { useEffect, useInsertionEffect, useMemo, useRef, useState } from "react";
 import { toast } from "sonner";
 import { z } from "zod";
 import { Button } from "../ui/button";
@@ -41,7 +41,12 @@ function Header() {
 		<div className="sticky top-0 flex items-start gap-2">
 			<div className="flex-1 rounded-xl border bg-fd-card p-3 text-fd-card-foreground">
 				<p className="mb-2 text-sm font-medium">Ask AI</p>
-				<p className="text-xs text-fd-muted-foreground">Powered by Gemma</p>
+				<p className="text-xs text-fd-muted-foreground">
+					Powered by{" "}
+					<a href="https://gemini.google.com" target="_blank" rel="noreferrer noopener">
+						Gemma
+					</a>
+				</p>
 			</div>
 
 			<Button
@@ -97,11 +102,13 @@ function SearchAIActions() {
 	);
 }
 
+const StorageKeyInput = "__ai_search_input";
+
 function SearchAIInput(props: InferProps<"form">) {
 	const { className, ...restOfProps } = props;
 	const { sendMessage, status, stop } = useChatContext();
 
-	const [input, setInput] = useState("");
+	const [input, setInput] = useState(() => localStorage.getItem(StorageKeyInput) ?? "");
 
 	const isLoading = status === "streaming" || status === "submitted";
 
@@ -111,6 +118,10 @@ function SearchAIInput(props: InferProps<"form">) {
 		setInput("");
 	};
 
+	useInsertionEffect(() => {
+		localStorage.setItem(StorageKeyInput, input);
+	}, [input]);
+
 	useEffect(() => {
 		if (isLoading) {
 			document.querySelector<HTMLElement>("#nd-ai-input")?.focus();
@@ -286,7 +297,9 @@ const roleName = {
 	user: "you",
 } satisfies Record<Exclude<UIMessage["role"], "system">, string> as Record<string, string>;
 
-export function AISearchTrigger() {
+export function AISearchRoot(props: { children: React.ReactNode }) {
+	const { children } = props;
+
 	const [open, setOpen] = useState(false);
 
 	const chat = useChat({
@@ -296,6 +309,36 @@ export function AISearchTrigger() {
 		}),
 	});
 
+	const contextValue = useMemo<GeneralContextType>(
+		() => ({ chat, open, setOpen }) satisfies GeneralContextType,
+		[chat, open]
+	);
+
+	return <GeneralContextProvider value={contextValue}>{children}</GeneralContextProvider>;
+}
+
+export function AISearchTrigger() {
+	const { open, setOpen } = useGeneralContext();
+
+	return (
+		<Button
+			theme="secondary"
+			className={cnMerge(
+				`fixed end-[calc(--spacing(4)+var(--removed-body-scroll-bar-size,0px))] bottom-4 z-20 w-24
+				gap-3 rounded-2xl text-fd-muted-foreground shadow-lg transition-[translate,opacity]`,
+				open && "translate-y-10 opacity-0"
+			)}
+			onClick={() => setOpen(true)}
+		>
+			<MessageCircleIcon className="size-4.5" />
+			Ask AI
+		</Button>
+	);
+}
+
+export function AISearchPanel() {
+	const { chat, open, setOpen } = useGeneralContext();
+
 	const onKeyPress = useCallbackRef((event: KeyboardEvent) => {
 		if (event.key === "Escape" && open) {
 			setOpen(false);
@@ -314,80 +357,87 @@ export function AISearchTrigger() {
 		return () => cleanup();
 	}, [onKeyPress]);
 
-	const contextValue = useMemo<GeneralContextType>(
-		() => ({ chat, open, setOpen }) satisfies GeneralContextType,
-		[chat, open]
-	);
-
 	return (
-		<GeneralContextProvider value={contextValue}>
+		<>
 			<style>
 				{css`
 					@keyframes ask-ai-open {
 						from {
-							translate: 100% 0;
+							width: 0px;
+						}
+						to {
+							width: var(--ai-chat-width);
 						}
 					}
-
 					@keyframes ask-ai-close {
+						from {
+							width: var(--ai-chat-width);
+						}
 						to {
-							translate: 100% 0;
-							opacity: 0;
+							width: 0px;
 						}
 					}
 				`}
 			</style>
 
+			<Presence present={open}>
+				<div
+					data-state={open ? "open" : "closed"}
+					className="fixed inset-0 z-30 bg-fd-overlay backdrop-blur-xs
+						data-[state=closed]:animate-fd-fade-out data-[state=open]:animate-fd-fade-in lg:hidden"
+					onClick={() => setOpen(false)}
+				/>
+			</Presence>
+
 			<Presence present={open}>
 				<div
 					className={cnMerge(
-						`fixed inset-y-2 z-30 flex flex-col rounded-2xl border bg-fd-popover p-2
-						text-fd-popover-foreground shadow-lg max-sm:inset-x-2 sm:end-2 sm:w-[460px]`,
-						open ? "animate-[ask-ai-open_300ms]" : "animate-[ask-ai-close_300ms]"
+						`z-30 overflow-hidden bg-fd-popover text-fd-popover-foreground [--ai-chat-width:400px]
+						xl:[--ai-chat-width:460px]`,
+						`max-lg:fixed max-lg:inset-x-2 max-lg:top-4 max-lg:rounded-2xl max-lg:border
+						max-lg:shadow-xl`,
+						`lg:sticky lg:top-0 lg:ms-auto lg:h-dvh lg:border-s
+						lg:in-[#nd-docs-layout]:[grid-area:toc] lg:in-[#nd-notebook-layout]:col-start-5
+						lg:in-[#nd-notebook-layout]:row-span-full`,
+						open ?
+							"animate-fd-dialog-in lg:animate-[ask-ai-open_200ms]"
+						:	"animate-fd-dialog-out lg:animate-[ask-ai-close_200ms]"
 					)}
 				>
-					<Header />
-					<MessageList
-						className="flex-1 overscroll-contain px-3 py-4"
-						style={{
-							maskImage:
-								"linear-gradient(to bottom, transparent, white 1rem, white calc(100% - 1rem), transparent 100%)",
-						}}
-					>
-						<div className="flex flex-col gap-4">
-							{chat.messages
-								.filter((msg) => msg.role !== "system")
-								.map((item) => (
-									<Message key={item.id} message={item} />
-								))}
-						</div>
-					</MessageList>
-
 					<div
-						className="rounded-xl border bg-fd-card text-fd-card-foreground has-focus-visible:ring-2
-							has-focus-visible:ring-fd-ring"
+						className="flex size-full flex-col p-2 max-lg:max-h-[80dvh] lg:w-(--ai-chat-width)
+							xl:p-4"
 					>
-						<SearchAIInput />
-						<div className="flex items-center gap-1.5 p-1 empty:hidden">
-							<SearchAIActions />
+						<Header />
+
+						<MessageList
+							className="flex-1 overscroll-contain px-3 py-4"
+							style={{
+								maskImage:
+									"linear-gradient(to bottom, transparent, white 1rem, white calc(100% - 1rem), transparent 100%)",
+							}}
+						>
+							<div className="flex flex-col gap-4">
+								{chat.messages
+									.filter((msg) => msg.role !== "system")
+									.map((item) => (
+										<Message key={item.id} message={item} />
+									))}
+							</div>
+						</MessageList>
+
+						<div
+							className="rounded-xl border bg-fd-card text-fd-card-foreground
+								has-focus-visible:ring-2 has-focus-visible:ring-fd-ring"
+						>
+							<SearchAIInput />
+							<div className="flex items-center gap-1.5 p-1 empty:hidden">
+								<SearchAIActions />
+							</div>
 						</div>
 					</div>
 				</div>
 			</Presence>
-
-			<Button
-				unstyled={true}
-				className={cnMerge(
-					`fixed bottom-4 z-20 flex h-10 w-24 items-center gap-3 rounded-2xl border bg-fd-secondary
-					px-2 text-sm font-medium text-fd-muted-foreground shadow-lg transition-[translate,opacity]`,
-					"end-[calc(var(--removed-body-scroll-bar-size,0px)+var(--fd-layout-offset)+1rem)]",
-					open && "translate-y-10 opacity-0"
-				)}
-				onClick={() => setOpen(true)}
-			>
-				<MessageCircleIcon className="size-4.5" />
-				Ask AI
-			</Button>
-		</GeneralContextProvider>
+		</>
 	);
 }