@tonyclaw/llm-inspector 1.16.2 → 1.16.4

package/src/components/ui/json-viewer.tsx

@@ -1,11 +1,15 @@
 import { Check, ChevronDown, ChevronRight, ChevronsDown, ChevronsUp, Copy } from "lucide-react";
-import { type JSX, memo, useState } from "react";
+import { type JSX, memo, useCallback, useMemo, useState, useTransition } from "react";
 import ReactMarkdown from "react-markdown";
 import { cn } from "../../lib/utils";
+import { Button } from "./button";
 import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "./tooltip";
 
 type JsonPrimitive = string | number | boolean | null;
-type JsonValue = JsonPrimitive | ReadonlyArray<JsonValue> | Readonly<{ [key: string]: JsonValue }>;
+export type JsonValue =
+  | JsonPrimitive
+  | ReadonlyArray<JsonValue>
+  | Readonly<{ [key: string]: JsonValue }>;
 
 type DataType = "string" | "number" | "boolean" | "null" | "array" | "object";
 
@@ -33,29 +37,141 @@ function isExpandable(value: JsonValue): boolean {
   return value !== null && (Array.isArray(value) || typeof value === "object");
 }
 
-function getPropertyValue(obj: object, key: string): unknown {
-  const descriptor = Object.getOwnPropertyDescriptor(obj, key);
-  if (descriptor === undefined) return undefined;
-  return descriptor.value;
-}
-
 function getEntries(value: JsonValue): ReadonlyArray<readonly [string, JsonValue]> {
   if (Array.isArray(value)) {
-    return value.map((item, index) => [String(index), item] as const);
+    return value.map((item, index): readonly [string, JsonValue] => [String(index), item]);
   }
   if (typeof value === "object" && value !== null) {
-    return Object.keys(value).map((key) => {
-      const raw = getPropertyValue(value, key);
-      return [key, safeJsonValue(raw)] as const;
-    });
+    return Object.entries(value);
   }
   return [];
 }
 
-function getItemCount(value: JsonValue): number {
-  if (Array.isArray(value)) return value.length;
-  if (typeof value === "object" && value !== null) return Object.keys(value).length;
-  return 0;
+export const JSON_FULL_EXPANSION_NODE_LIMIT = 1000;
+export const JSON_LARGE_EXPANSION_DEPTH = 2;
+
+export type JsonExpansionPolicy = {
+  depth: number;
+  limited: boolean;
+  countedNodes: number;
+};
+
+export function countJsonNodes(value: JsonValue, limit = Number.POSITIVE_INFINITY): number {
+  let count = 0;
+  const pending: JsonValue[] = [value];
+
+  while (pending.length > 0) {
+    const current = pending.pop();
+    if (current === undefined) continue;
+    count += 1;
+    if (count > limit) return count;
+    if (!isExpandable(current)) continue;
+    for (const [, child] of getEntries(current)) {
+      pending.push(child);
+    }
+  }
+
+  return count;
+}
+
+export function getJsonExpansionPolicy(value: JsonValue): JsonExpansionPolicy {
+  const countedNodes = countJsonNodes(value, JSON_FULL_EXPANSION_NODE_LIMIT);
+  const limited = countedNodes > JSON_FULL_EXPANSION_NODE_LIMIT;
+  return {
+    depth: limited ? JSON_LARGE_EXPANSION_DEPTH : Number.POSITIVE_INFINITY,
+    limited,
+    countedNodes,
+  };
+}
+
+export type JsonBulkExpansion = {
+  parsedData: JsonValue | null;
+  policy: JsonExpansionPolicy | null;
+  isExpanded: boolean;
+  toggle: () => void;
+  isPending: boolean;
+  bulkDepth: number;
+  bulkRevision: number;
+};
+
+export function useJsonBulkExpansion(text: string | null): JsonBulkExpansion {
+  const parsed = useMemo(() => (text === null ? null : parseJsonText(text)), [text]);
+  const parsedData = parsed?.kind === "json" ? parsed.data : null;
+  const policy = useMemo(
+    () => (parsedData === null ? null : getJsonExpansionPolicy(parsedData)),
+    [parsedData],
+  );
+  const [isExpanded, setIsExpanded] = useState(false);
+  const [bulkDepth, setBulkDepth] = useState(0);
+  const [bulkRevision, setBulkRevision] = useState(0);
+  const [isPending, startTransition] = useTransition();
+
+  const toggle = useCallback(() => {
+    const nextExpanded = !isExpanded;
+    const targetDepth = nextExpanded && policy !== null ? policy.depth : 0;
+    startTransition(() => {
+      setIsExpanded(nextExpanded);
+      setBulkDepth(targetDepth);
+      setBulkRevision((current) => current + 1);
+    });
+  }, [isExpanded, policy]);
+
+  return { parsedData, policy, isExpanded, toggle, isPending, bulkDepth, bulkRevision };
+}
+
+export function JsonExpansionButton({
+  policy,
+  isExpanded,
+  isPending,
+  onToggle,
+}: {
+  policy: JsonExpansionPolicy | null;
+  isExpanded: boolean;
+  isPending: boolean;
+  onToggle: () => void;
+}): JSX.Element | null {
+  if (policy === null) return null;
+
+  const label = isPending
+    ? "Updating..."
+    : isExpanded
+      ? "Collapse all"
+      : policy.limited
+        ? `Expand ${JSON_LARGE_EXPANSION_DEPTH} levels`
+        : "Expand all";
+
+  const tooltip =
+    policy.limited && !isExpanded
+      ? `Large JSON detected; expand ${JSON_LARGE_EXPANSION_DEPTH} levels to protect rendering performance`
+      : isExpanded
+        ? "Collapse all JSON nodes"
+        : "Expand all JSON nodes";
+
+  return (
+    <Tooltip>
+      <TooltipTrigger asChild>
+        <Button
+          variant="outline"
+          size="sm"
+          className="h-7 text-xs"
+          onClick={(e) => {
+            e.stopPropagation();
+            onToggle();
+          }}
+          disabled={isPending}
+          aria-pressed={isExpanded}
+        >
+          {isExpanded ? (
+            <ChevronsUp className="size-3 mr-1" />
+          ) : (
+            <ChevronsDown className="size-3 mr-1" />
+          )}
+          {label}
+        </Button>
+      </TooltipTrigger>
+      <TooltipContent>{tooltip}</TooltipContent>
+    </Tooltip>
+  );
 }
 
 const STRING_TRUNCATE_LIMIT = 120;
@@ -206,16 +322,9 @@ type JsonNodeProps = {
   isArrayItem: boolean;
 };
 
-function hasExpandableDescendant(value: JsonValue): boolean {
-  if (!isExpandable(value)) return false;
-  return getEntries(value).some(([, v]) => isExpandable(v));
-}
-
 function ExpandCollapseButton({
-  allExpanded,
   onClick,
 }: {
-  allExpanded: boolean;
   onClick: (e: React.MouseEvent) => void;
 }): JSX.Element {
   return (
@@ -223,13 +332,9 @@ function ExpandCollapseButton({
       type="button"
       onClick={onClick}
       className="opacity-0 group-hover/row:opacity-100 hover:bg-muted p-0.5 rounded transition-opacity shrink-0"
-      title={allExpanded ? "Collapse all" : "Expand all"}
+      title="Expand all descendants"
     >
-      {allExpanded ? (
-        <ChevronsUp className="size-3.5 text-muted-foreground" />
-      ) : (
-        <ChevronsDown className="size-3.5 text-muted-foreground" />
-      )}
+      <ChevronsDown className="size-3.5 text-muted-foreground" />
     </button>
   );
 }
@@ -244,30 +349,41 @@ const JsonNode = memo(function JsonNode({
   const [expanded, setExpanded] = useState(level < defaultExpandDepth);
   const [childResetKey, setChildResetKey] = useState(0);
   const [childDepthOverride, setChildDepthOverride] = useState<number | null>(null);
-  const [allExpanded, setAllExpanded] = useState(false);
   const expandable = isExpandable(value);
+  const fullyExpanded = childDepthOverride === Number.POSITIVE_INFINITY;
+  const entries = useMemo(() => getEntries(value), [value]);
+  const hasExpandableChild = useMemo(
+    () => entries.some(([, child]) => isExpandable(child)),
+    [entries],
+  );
 
   const dataType = classifyValue(value);
   const openBracket = dataType === "array" ? "[" : "{";
   const closeBracket = dataType === "array" ? "]" : "}";
 
-  function toggleDeepExpansion(): void {
-    if (allExpanded) {
-      setExpanded(false);
-      setChildDepthOverride(0);
-      setChildResetKey((k) => k + 1);
-      setAllExpanded(false);
+  function expandAllDeep(): void {
+    setExpanded(true);
+    setChildDepthOverride(Number.POSITIVE_INFINITY);
+    setChildResetKey((k) => k + 1);
+  }
+
+  function collapseToDefault(): void {
+    setExpanded(false);
+    setChildDepthOverride(0);
+    setChildResetKey((k) => k + 1);
+  }
+
+  function handleRowToggle(): void {
+    if (fullyExpanded) {
+      collapseToDefault();
     } else {
-      setExpanded(true);
-      setChildDepthOverride(Number.POSITIVE_INFINITY);
-      setChildResetKey((k) => k + 1);
-      setAllExpanded(true);
+      setExpanded(!expanded);
     }
   }
 
   function handleExpandAll(e: React.MouseEvent): void {
     e.stopPropagation();
-    toggleDeepExpansion();
+    expandAllDeep();
   }
 
   const effectiveChildDepth = childDepthOverride ?? defaultExpandDepth;
@@ -279,29 +395,18 @@ const JsonNode = memo(function JsonNode({
           "flex items-start gap-1 py-0.5 px-1 -ml-1 rounded-sm group/row",
           expandable && "cursor-pointer hover:bg-muted/50",
         )}
-        onClick={
-          expandable
-            ? () => {
-                if (expanded) {
-                  setAllExpanded(false);
-                }
-                setExpanded(!expanded);
-              }
-            : undefined
-        }
+        onClick={expandable ? handleRowToggle : undefined}
         onKeyDown={
           expandable
             ? (e) => {
                 if (e.key === "Enter" || e.key === " ") {
                   e.preventDefault();
-                  setExpanded(!expanded);
+                  handleRowToggle();
                 }
               }
             : undefined
         }
-        onDoubleClick={
-          expandable && hasExpandableDescendant(value) ? () => toggleDeepExpansion() : undefined
-        }
+        onDoubleClick={expandable && hasExpandableChild ? () => expandAllDeep() : undefined}
         role={expandable ? "button" : undefined}
         tabIndex={expandable ? 0 : undefined}
       >
@@ -328,7 +433,7 @@ const JsonNode = memo(function JsonNode({
             {openBracket}
             <span className="text-muted-foreground/60 text-xs">
               {" "}
-              {getItemCount(value)} {getItemCount(value) === 1 ? "item" : "items"} {closeBracket}
+              {entries.length} {entries.length === 1 ? "item" : "items"} {closeBracket}
             </span>
           </span>
         ) : (
@@ -337,15 +442,15 @@ const JsonNode = memo(function JsonNode({
           </span>
         )}
 
-        {expandable && hasExpandableDescendant(value) && (
-          <ExpandCollapseButton allExpanded={allExpanded} onClick={handleExpandAll} />
+        {expandable && hasExpandableChild && !fullyExpanded && (
+          <ExpandCollapseButton onClick={handleExpandAll} />
         )}
         <CopyButton value={value} />
       </div>
 
       {expandable && expanded && (
         <div className="pl-4" key={childResetKey}>
-          {getEntries(value).map(([key, childValue]) => (
+          {entries.map(([key, childValue]) => (
             <JsonNode
               key={key}
               name={key}
@@ -367,15 +472,23 @@ export type JsonViewerProps = {
   defaultExpandDepth?: number;
   className?: string;
   showCopy?: boolean;
+  bulkDepth?: number;
+  bulkRevision?: number;
 };
 
 export function JsonViewer({
   data,
-  defaultExpandDepth = 2,
+  defaultExpandDepth = 0,
   className,
   showCopy = false,
+  bulkDepth: controlledBulkDepth,
+  bulkRevision: controlledBulkRevision,
 }: JsonViewerProps): JSX.Element {
   const expandable = isExpandable(data);
+  const entries = useMemo(() => getEntries(data), [data]);
+
+  const bulkDepth = controlledBulkDepth ?? defaultExpandDepth;
+  const bulkRevision = controlledBulkRevision ?? 0;
 
   if (!expandable) {
     return (
@@ -396,17 +509,23 @@ export function JsonViewer({
   return (
     <TooltipProvider>
       <div className={cn("font-mono text-xs leading-relaxed", className)}>
-        {showCopy && <CopyValueButton value={data} />}
-        {getEntries(data).map(([key, childValue]) => (
-          <JsonNode
-            key={key}
-            name={key}
-            value={childValue}
-            level={0}
-            defaultExpandDepth={defaultExpandDepth}
-            isArrayItem={isArray}
-          />
-        ))}
+        {showCopy && (
+          <div className="mb-2 flex items-center justify-end gap-2">
+            <CopyValueButton value={data} />
+          </div>
+        )}
+        <div key={bulkRevision}>
+          {entries.map(([key, childValue]) => (
+            <JsonNode
+              key={key}
+              name={key}
+              value={childValue}
+              level={0}
+              defaultExpandDepth={bulkDepth}
+              isArrayItem={isArray}
+            />
+          ))}
+        </div>
       </div>
     </TooltipProvider>
   );
@@ -418,28 +537,45 @@ export type JsonViewerFromStringProps = {
   className?: string;
 };
 
-export function JsonViewerFromString({
+type ParsedJsonText =
+  | {
+      kind: "json";
+      data: JsonValue;
+    }
+  | {
+      kind: "text";
+    };
+
+function parseJsonText(text: string): ParsedJsonText {
+  try {
+    const parsed: unknown = JSON.parse(text);
+    return { kind: "json", data: safeJsonValue(parsed) };
+  } catch {
+    return { kind: "text" };
+  }
+}
+
+export const JsonViewerFromString = memo(function JsonViewerFromString({
   text,
-  defaultExpandDepth = 2,
+  defaultExpandDepth = 0,
   className,
 }: JsonViewerFromStringProps): JSX.Element {
-  try {
-    const parsed: unknown = JSON.parse(text);
+  const parsed = useMemo(() => parseJsonText(text), [text]);
+
+  if (parsed.kind === "json") {
     return (
       <JsonViewer
-        data={safeJsonValue(parsed)}
+        data={parsed.data}
         defaultExpandDepth={defaultExpandDepth}
         className={className}
       />
     );
-  } catch {
-    return (
-      <pre className={cn("font-mono text-xs whitespace-pre-wrap break-words", className)}>
-        {text}
-      </pre>
-    );
   }
-}
+
+  return (
+    <pre className={cn("font-mono text-xs whitespace-pre-wrap break-words", className)}>{text}</pre>
+  );
+});
 
 export function safeJsonValue(value: unknown): JsonValue {
   if (value === null || value === undefined) return null;