work-agent 0.1.0

package/components/ai-elements/jsx-preview.tsx

@@ -0,0 +1,250 @@
+"use client";
+
+import type { ComponentProps, ReactNode } from "react";
+import type { TProps as JsxParserProps } from "react-jsx-parser";
+
+import { cn } from "@/lib/utils";
+import { AlertCircle } from "lucide-react";
+import {
+  createContext,
+  memo,
+  useCallback,
+  useContext,
+  useEffect,
+  useMemo,
+  useRef,
+  useState,
+} from "react";
+import JsxParser from "react-jsx-parser";
+
+interface JSXPreviewContextValue {
+  jsx: string;
+  processedJsx: string;
+  error: Error | null;
+  setError: (error: Error | null) => void;
+  components: JsxParserProps["components"];
+  bindings: JsxParserProps["bindings"];
+  onErrorProp?: (error: Error) => void;
+}
+
+const JSXPreviewContext = createContext<JSXPreviewContextValue | null>(null);
+
+const TAG_REGEX = /<\/?([a-zA-Z][a-zA-Z0-9]*)\s*([^>]*?)(\/)?>/;
+
+export const useJSXPreview = () => {
+  const context = useContext(JSXPreviewContext);
+  if (!context) {
+    throw new Error("JSXPreview components must be used within JSXPreview");
+  }
+  return context;
+};
+
+const matchJsxTag = (code: string) => {
+  if (code.trim() === "") {
+    return null;
+  }
+
+  const match = code.match(TAG_REGEX);
+
+  if (!match || match.index === undefined) {
+    return null;
+  }
+
+  const [fullMatch, tagName, attributes, selfClosing] = match;
+
+  let type: "self-closing" | "closing" | "opening";
+  if (selfClosing) {
+    type = "self-closing";
+  } else if (fullMatch.startsWith("</")) {
+    type = "closing";
+  } else {
+    type = "opening";
+  }
+
+  return {
+    attributes: attributes.trim(),
+    endIndex: match.index + fullMatch.length,
+    startIndex: match.index,
+    tag: fullMatch,
+    tagName,
+    type,
+  };
+};
+
+const completeJsxTag = (code: string) => {
+  const stack: string[] = [];
+  let result = "";
+  let currentPosition = 0;
+
+  while (currentPosition < code.length) {
+    const match = matchJsxTag(code.slice(currentPosition));
+    if (!match) {
+      // No more tags found, append remaining content
+      result += code.slice(currentPosition);
+      break;
+    }
+    const { tagName, type, endIndex } = match;
+
+    // Include any text content before this tag
+    result += code.slice(currentPosition, currentPosition + endIndex);
+
+    if (type === "opening") {
+      stack.push(tagName);
+    } else if (type === "closing") {
+      stack.pop();
+    }
+
+    currentPosition += endIndex;
+  }
+
+  return (
+    result +
+    stack
+      .toReversed()
+      .map((tag) => `</${tag}>`)
+      .join("")
+  );
+};
+
+export type JSXPreviewProps = ComponentProps<"div"> & {
+  jsx: string;
+  isStreaming?: boolean;
+  components?: JsxParserProps["components"];
+  bindings?: JsxParserProps["bindings"];
+  onError?: (error: Error) => void;
+};
+
+export const JSXPreview = memo(
+  ({
+    jsx,
+    isStreaming = false,
+    components,
+    bindings,
+    onError,
+    className,
+    children,
+    ...props
+  }: JSXPreviewProps) => {
+    const [prevJsx, setPrevJsx] = useState(jsx);
+    const [error, setError] = useState<Error | null>(null);
+
+    // Clear error when jsx changes (derived state pattern)
+    if (jsx !== prevJsx) {
+      setPrevJsx(jsx);
+      setError(null);
+    }
+
+    const processedJsx = useMemo(
+      () => (isStreaming ? completeJsxTag(jsx) : jsx),
+      [jsx, isStreaming]
+    );
+
+    return (
+      <JSXPreviewContext.Provider
+        value={{
+          bindings,
+          components,
+          error,
+          jsx,
+          onErrorProp: onError,
+          processedJsx,
+          setError,
+        }}
+      >
+        <div className={cn("relative", className)} {...props}>
+          {children}
+        </div>
+      </JSXPreviewContext.Provider>
+    );
+  }
+);
+
+JSXPreview.displayName = "JSXPreview";
+
+export type JSXPreviewContentProps = Omit<ComponentProps<"div">, "children">;
+
+export const JSXPreviewContent = memo(
+  ({ className, ...props }: JSXPreviewContentProps) => {
+    const { processedJsx, components, bindings, setError, onErrorProp } =
+      useJSXPreview();
+    const errorReportedRef = useRef<string | null>(null);
+
+    // Reset error tracking when jsx changes
+    // biome-ignore lint/correctness/useExhaustiveDependencies: processedJsx change should reset tracking
+    useEffect(() => {
+      errorReportedRef.current = null;
+    }, [processedJsx]);
+
+    const handleError = useCallback(
+      (err: Error) => {
+        // Prevent duplicate error reports for the same jsx
+        if (errorReportedRef.current === processedJsx) {
+          return;
+        }
+        errorReportedRef.current = processedJsx;
+        setError(err);
+        onErrorProp?.(err);
+      },
+      [processedJsx, onErrorProp, setError]
+    );
+
+    return (
+      <div className={cn("jsx-preview-content", className)} {...props}>
+        <JsxParser
+          bindings={bindings}
+          components={components}
+          jsx={processedJsx}
+          onError={handleError}
+          renderInWrapper={false}
+        />
+      </div>
+    );
+  }
+);
+
+JSXPreviewContent.displayName = "JSXPreviewContent";
+
+export type JSXPreviewErrorProps = ComponentProps<"div"> & {
+  children?: ReactNode | ((error: Error) => ReactNode);
+};
+
+const renderChildren = (
+  children: ReactNode | ((error: Error) => ReactNode),
+  error: Error
+): ReactNode => {
+  if (typeof children === "function") {
+    return children(error);
+  }
+  return children;
+};
+
+export const JSXPreviewError = memo(
+  ({ className, children, ...props }: JSXPreviewErrorProps) => {
+    const { error } = useJSXPreview();
+
+    if (!error) {
+      return null;
+    }
+
+    return (
+      <div
+        className={cn(
+          "flex items-center gap-2 rounded-md border border-destructive/50 bg-destructive/10 p-3 text-destructive text-sm",
+          className
+        )}
+        {...props}
+      >
+        {children ? (
+          renderChildren(children, error)
+        ) : (
+          <>
+            <AlertCircle className="size-4 shrink-0" />
+            <span>{error.message}</span>
+          </>
+        )}
+      </div>
+    );
+  }
+);
+
+JSXPreviewError.displayName = "JSXPreviewError";

package/components/ai-elements/message.tsx

@@ -0,0 +1,367 @@
+"use client";
+
+import type { UIMessage } from "ai";
+import type { ComponentProps, HTMLAttributes, ReactElement } from "react";
+
+import { Button } from "@/components/ui/button";
+import {
+  ButtonGroup,
+  ButtonGroupText,
+} from "@/components/ui/button-group";
+import {
+  Tooltip,
+  TooltipContent,
+  TooltipProvider,
+  TooltipTrigger,
+} from "@/components/ui/tooltip";
+import { cn } from "@/lib/utils";
+import { cjk } from "@streamdown/cjk";
+import { code } from "@streamdown/code";
+import { math } from "@streamdown/math";
+import { mermaid } from "@streamdown/mermaid";
+import { ChevronLeftIcon, ChevronRightIcon } from "lucide-react";
+import {
+  createContext,
+  memo,
+  useCallback,
+  useContext,
+  useEffect,
+  useMemo,
+  useState,
+} from "react";
+import { Streamdown } from "streamdown";
+
+export type MessageProps = HTMLAttributes<HTMLDivElement> & {
+  from: UIMessage["role"];
+};
+
+export const Message = ({ className, from, ...props }: MessageProps) => (
+  <div
+    className={cn(
+      "group flex w-full max-w-[95%] flex-col gap-2",
+      from === "user" ? "is-user ml-auto justify-end" : "is-assistant",
+      className
+    )}
+    {...props}
+  />
+);
+
+export type MessageContentProps = HTMLAttributes<HTMLDivElement>;
+
+export const MessageContent = ({
+  children,
+  className,
+  ...props
+}: MessageContentProps) => (
+  <div
+    className={cn(
+      "is-user:dark flex w-fit min-w-0 max-w-full flex-col gap-2 overflow-hidden text-sm",
+      "group-[.is-user]:ml-auto group-[.is-user]:rounded-lg group-[.is-user]:bg-secondary group-[.is-user]:px-4 group-[.is-user]:py-3 group-[.is-user]:text-foreground",
+      "group-[.is-assistant]:text-foreground",
+      className
+    )}
+    {...props}
+  >
+    {children}
+  </div>
+);
+
+export type MessageActionsProps = ComponentProps<"div">;
+
+export const MessageActions = ({
+  className,
+  children,
+  ...props
+}: MessageActionsProps) => (
+  <div className={cn("flex items-center gap-1", className)} {...props}>
+    {children}
+  </div>
+);
+
+export type MessageActionProps = ComponentProps<typeof Button> & {
+  tooltip?: string;
+  label?: string;
+};
+
+export const MessageAction = ({
+  tooltip,
+  children,
+  label,
+  variant = "ghost",
+  size = "icon-sm",
+  ...props
+}: MessageActionProps) => {
+  const button = (
+    <Button size={size} type="button" variant={variant} {...props}>
+      {children}
+      <span className="sr-only">{label || tooltip}</span>
+    </Button>
+  );
+
+  if (tooltip) {
+    return (
+      <TooltipProvider>
+        <Tooltip>
+          <TooltipTrigger asChild>{button}</TooltipTrigger>
+          <TooltipContent>
+            <p>{tooltip}</p>
+          </TooltipContent>
+        </Tooltip>
+      </TooltipProvider>
+    );
+  }
+
+  return button;
+};
+
+interface MessageBranchContextType {
+  currentBranch: number;
+  totalBranches: number;
+  goToPrevious: () => void;
+  goToNext: () => void;
+  branches: ReactElement[];
+  setBranches: (branches: ReactElement[]) => void;
+}
+
+const MessageBranchContext = createContext<MessageBranchContextType | null>(
+  null
+);
+
+const useMessageBranch = () => {
+  const context = useContext(MessageBranchContext);
+
+  if (!context) {
+    throw new Error(
+      "MessageBranch components must be used within MessageBranch"
+    );
+  }
+
+  return context;
+};
+
+export type MessageBranchProps = HTMLAttributes<HTMLDivElement> & {
+  defaultBranch?: number;
+  onBranchChange?: (branchIndex: number) => void;
+};
+
+export const MessageBranch = ({
+  defaultBranch = 0,
+  onBranchChange,
+  className,
+  ...props
+}: MessageBranchProps) => {
+  const [currentBranch, setCurrentBranch] = useState(defaultBranch);
+  const [branches, setBranches] = useState<ReactElement[]>([]);
+
+  const handleBranchChange = useCallback(
+    (newBranch: number) => {
+      setCurrentBranch(newBranch);
+      onBranchChange?.(newBranch);
+    },
+    [onBranchChange]
+  );
+
+  const goToPrevious = useCallback(() => {
+    const newBranch =
+      currentBranch > 0 ? currentBranch - 1 : branches.length - 1;
+    handleBranchChange(newBranch);
+  }, [currentBranch, branches.length, handleBranchChange]);
+
+  const goToNext = useCallback(() => {
+    const newBranch =
+      currentBranch < branches.length - 1 ? currentBranch + 1 : 0;
+    handleBranchChange(newBranch);
+  }, [currentBranch, branches.length, handleBranchChange]);
+
+  const contextValue = useMemo<MessageBranchContextType>(
+    () => ({
+      branches,
+      currentBranch,
+      goToNext,
+      goToPrevious,
+      setBranches,
+      totalBranches: branches.length,
+    }),
+    [branches, currentBranch, goToNext, goToPrevious]
+  );
+
+  return (
+    <MessageBranchContext.Provider value={contextValue}>
+      <div
+        className={cn("grid w-full gap-2 [&>div]:pb-0", className)}
+        {...props}
+      />
+    </MessageBranchContext.Provider>
+  );
+};
+
+export type MessageBranchContentProps = HTMLAttributes<HTMLDivElement>;
+
+export const MessageBranchContent = ({
+  children,
+  ...props
+}: MessageBranchContentProps) => {
+  const { currentBranch, setBranches, branches } = useMessageBranch();
+  const childrenArray = useMemo(
+    () => (Array.isArray(children) ? children : [children]),
+    [children]
+  );
+
+  // Use useEffect to update branches when they change
+  useEffect(() => {
+    if (branches.length !== childrenArray.length) {
+      setBranches(childrenArray);
+    }
+  }, [childrenArray, branches, setBranches]);
+
+  return childrenArray.map((branch, index) => (
+    <div
+      className={cn(
+        "grid gap-2 overflow-hidden [&>div]:pb-0",
+        index === currentBranch ? "block" : "hidden"
+      )}
+      key={branch.key}
+      {...props}
+    >
+      {branch}
+    </div>
+  ));
+};
+
+export type MessageBranchSelectorProps = ComponentProps<typeof ButtonGroup>;
+
+export const MessageBranchSelector = ({
+  className,
+  ...props
+}: MessageBranchSelectorProps) => {
+  const { totalBranches } = useMessageBranch();
+
+  // Don't render if there's only one branch
+  if (totalBranches <= 1) {
+    return null;
+  }
+
+  return (
+    <ButtonGroup
+      className={cn(
+        "[&>*:not(:first-child)]:rounded-l-md [&>*:not(:last-child)]:rounded-r-md",
+        className
+      )}
+      orientation="horizontal"
+      {...props}
+    />
+  );
+};
+
+export type MessageBranchPreviousProps = ComponentProps<typeof Button>;
+
+export const MessageBranchPrevious = ({
+  children,
+  ...props
+}: MessageBranchPreviousProps) => {
+  const { goToPrevious, totalBranches } = useMessageBranch();
+
+  return (
+    <Button
+      aria-label="Previous branch"
+      disabled={totalBranches <= 1}
+      onClick={goToPrevious}
+      size="icon-sm"
+      type="button"
+      variant="ghost"
+      {...props}
+    >
+      {children ?? <ChevronLeftIcon size={14} />}
+    </Button>
+  );
+};
+
+export type MessageBranchNextProps = ComponentProps<typeof Button>;
+
+export const MessageBranchNext = ({
+  children,
+  ...props
+}: MessageBranchNextProps) => {
+  const { goToNext, totalBranches } = useMessageBranch();
+
+  return (
+    <Button
+      aria-label="Next branch"
+      disabled={totalBranches <= 1}
+      onClick={goToNext}
+      size="icon-sm"
+      type="button"
+      variant="ghost"
+      {...props}
+    >
+      {children ?? <ChevronRightIcon size={14} />}
+    </Button>
+  );
+};
+
+export type MessageBranchPageProps = HTMLAttributes<HTMLSpanElement>;
+
+export const MessageBranchPage = ({
+  className,
+  ...props
+}: MessageBranchPageProps) => {
+  const { currentBranch, totalBranches } = useMessageBranch();
+
+  return (
+    <ButtonGroupText
+      className={cn(
+        "border-none bg-transparent text-muted-foreground shadow-none",
+        className
+      )}
+      {...props}
+    >
+      {currentBranch + 1} of {totalBranches}
+    </ButtonGroupText>
+  );
+};
+
+export type MessageResponseProps = ComponentProps<typeof Streamdown>;
+
+export const MessageResponse = memo(
+  ({ className, ...props }: MessageResponseProps) => (
+    <Streamdown
+      className={cn(
+        "size-full [&>*:first-child]:mt-0 [&>*:last-child]:mb-0",
+        className
+      )}
+      plugins={{ code, cjk, math, mermaid }}
+      controls={{
+        code: true,
+        table: true,
+        mermaid: {
+          download: true,
+          copy: true,
+          fullscreen: true,
+          panZoom: true,
+        },
+      }}
+      {...props}
+    />
+  ),
+  (prevProps, nextProps) => prevProps.children === nextProps.children
+);
+
+MessageResponse.displayName = "MessageResponse";
+
+export type MessageToolbarProps = ComponentProps<"div">;
+
+export const MessageToolbar = ({
+  className,
+  children,
+  ...props
+}: MessageToolbarProps) => (
+  <div
+    className={cn(
+      "mt-4 flex w-full items-center justify-between gap-4",
+      className
+    )}
+    {...props}
+  >
+    {children}
+  </div>
+);