@gram-ai/elements 1.25.1 → 1.26.0

package/src/contexts/ElementsProvider.tsx

@@ -59,6 +59,7 @@ import {
   useConnectionStatusOptional,
 } from './ConnectionStatusContext'
 import { ToolExecutionProvider } from './ToolExecutionContext'
+import { ChatIdContext } from './ChatIdContext'
 
 /**
  * Extracts executable tools from frontend tool definitions.
@@ -93,14 +94,25 @@ export interface ElementsProviderProps {
   config: ElementsConfig
 }
 
-const BASE_SYSTEM_PROMPT = `You are a helpful assistant that can answer questions and help with tasks.`
+const BASE_SYSTEM_PROMPT = `You are a helpful assistant that can answer questions and help with tasks.
+
+Tool Result Display:
+Some tools have custom visual components that automatically render their results (you'll see a rich card/widget appear). For these, do not repeat the data - just add brief context or a follow-up question if needed.
+
+For tools WITHOUT custom components, you should present the data clearly - either as plain text for simple results, or using the UI code block format for structured data like lists of items, categories, or dashboards.`
 
 function mergeInternalSystemPromptWith(
   userSystemPrompt: string | undefined,
-  plugins: Plugin[]
+  plugins: Plugin[],
+  toolsWithCustomComponents: string[]
 ) {
+  const customToolsSection =
+    toolsWithCustomComponents.length > 0
+      ? `\n\nTools with custom visual components (DO NOT render UI widgets for these - they already display rich visuals):\n${toolsWithCustomComponents.map((t) => `- ${t}`).join('\n')}`
+      : ''
+
   return `
-  ${BASE_SYSTEM_PROMPT}
+  ${BASE_SYSTEM_PROMPT}${customToolsSection}
 
   User-provided System Prompt:
   ${userSystemPrompt ?? 'None provided'}
@@ -158,9 +170,13 @@ const ElementsProviderInner = ({ children, config }: ElementsProviderProps) => {
 
   const plugins = config.plugins ?? recommended
 
+  // Get list of tools that have custom components registered
+  const toolsWithCustomComponents = Object.keys(config.tools?.components ?? {})
+
   const systemPrompt = mergeInternalSystemPromptWith(
     config.systemPrompt,
-    plugins
+    plugins,
+    toolsWithCustomComponents
   )
 
   // Initialize error tracking on mount
@@ -176,6 +192,9 @@ const ElementsProviderInner = ({ children, config }: ElementsProviderProps) => {
   // When history is enabled, the thread adapter manages chat IDs instead
   const chatIdRef = useRef<string | null>(null)
 
+  // State to expose the current chat ID via context
+  const [currentChatId, setCurrentChatId] = useState<string | null>(null)
+
   const { data: mcpTools, mcpHeaders } = useMCPTools({
     auth,
     mcp: config.mcp,
@@ -288,6 +307,8 @@ const ElementsProviderInner = ({ children, config }: ElementsProviderProps) => {
         // chat ID on subsequent tool call requests.
         if (chatId) {
           mcpHeaders['Gram-Chat-ID'] = chatId
+          // Update the context state so consumers can access the current chat ID
+          setCurrentChatId(chatId)
         }
 
         const context = runtimeRef.current?.thread.getModelContext()
@@ -467,6 +488,8 @@ const ElementsProviderInner = ({ children, config }: ElementsProviderProps) => {
         localIdToUuidMap={localIdToUuidMapRef.current}
         currentRemoteIdRef={currentRemoteIdRef}
         executableTools={executableTools}
+        currentChatId={currentChatId}
+        setCurrentChatId={setCurrentChatId}
       >
         {children}
       </ElementsProviderWithHistory>
@@ -480,6 +503,7 @@ const ElementsProviderInner = ({ children, config }: ElementsProviderProps) => {
       runtimeRef={runtimeRef}
       frontendTools={frontendTools}
       executableTools={executableTools}
+      currentChatId={currentChatId}
     >
       {children}
     </ElementsProviderWithoutHistory>
@@ -506,23 +530,28 @@ interface ElementsProviderWithHistoryProps {
   localIdToUuidMap: Map<string, string>
   currentRemoteIdRef: React.RefObject<string | null>
   executableTools: ExecutableToolSet
+  currentChatId: string | null
+  setCurrentChatId: (chatId: string | null) => void
 }
 
 /**
- * Component that syncs the current thread's remoteId to a ref.
+ * Component that syncs the current thread's remoteId to a ref and updates the chat ID context.
  * Must be rendered inside AssistantRuntimeProvider to access the state.
  */
 const ThreadIdSync = ({
   remoteIdRef,
+  onChatIdChange,
 }: {
   remoteIdRef: React.RefObject<string | null>
+  onChatIdChange: (chatId: string | null) => void
 }) => {
   const remoteId = useAssistantState(
     ({ threadListItem }) => threadListItem.remoteId ?? null
   )
   useEffect(() => {
     remoteIdRef.current = remoteId
-  }, [remoteId, remoteIdRef])
+    onChatIdChange(remoteId)
+  }, [remoteId, remoteIdRef, onChatIdChange])
   return null
 }
 
@@ -537,6 +566,8 @@ const ElementsProviderWithHistory = ({
   localIdToUuidMap,
   currentRemoteIdRef,
   executableTools,
+  currentChatId,
+  setCurrentChatId,
 }: ElementsProviderWithHistoryProps) => {
   const threadListAdapter = useGramThreadListAdapter({
     apiUrl,
@@ -582,23 +613,28 @@ const ElementsProviderWithHistory = ({
 
   return (
     <AssistantRuntimeProvider runtime={runtime}>
-      <ThreadIdSync remoteIdRef={currentRemoteIdRef} />
+      <ThreadIdSync
+        remoteIdRef={currentRemoteIdRef}
+        onChatIdChange={setCurrentChatId}
+      />
       <HistoryProvider>
-        <ElementsContext.Provider value={contextValue}>
-          <ToolExecutionProvider tools={executableTools}>
-            <div
-              className={cn(
-                ROOT_SELECTOR,
-                (contextValue?.config.variant === 'standalone' ||
-                  contextValue?.config.variant === 'sidecar') &&
-                  'h-full'
-              )}
-            >
-              {children}
-            </div>
-            <FrontendTools tools={frontendTools} />
-          </ToolExecutionProvider>
-        </ElementsContext.Provider>
+        <ChatIdContext.Provider value={{ chatId: currentChatId }}>
+          <ElementsContext.Provider value={contextValue}>
+            <ToolExecutionProvider tools={executableTools}>
+              <div
+                className={cn(
+                  ROOT_SELECTOR,
+                  (contextValue?.config.variant === 'standalone' ||
+                    contextValue?.config.variant === 'sidecar') &&
+                    'h-full'
+                )}
+              >
+                {children}
+              </div>
+              <FrontendTools tools={frontendTools} />
+            </ToolExecutionProvider>
+          </ElementsContext.Provider>
+        </ChatIdContext.Provider>
       </HistoryProvider>
     </AssistantRuntimeProvider>
   )
@@ -613,6 +649,7 @@ interface ElementsProviderWithoutHistoryProps {
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   frontendTools: Record<string, AssistantTool | FrontendTool<any, any>>
   executableTools: ExecutableToolSet
+  currentChatId: string | null
 }
 
 const ElementsProviderWithoutHistory = ({
@@ -622,6 +659,7 @@ const ElementsProviderWithoutHistory = ({
   runtimeRef,
   frontendTools,
   executableTools,
+  currentChatId,
 }: ElementsProviderWithoutHistoryProps) => {
   const runtime = useChatRuntime({ transport })
 
@@ -632,21 +670,23 @@ const ElementsProviderWithoutHistory = ({
 
   return (
     <AssistantRuntimeProvider runtime={runtime}>
-      <ElementsContext.Provider value={contextValue}>
-        <ToolExecutionProvider tools={executableTools}>
-          <div
-            className={cn(
-              ROOT_SELECTOR,
-              (contextValue?.config.variant === 'standalone' ||
-                contextValue?.config.variant === 'sidecar') &&
-                'h-full'
-            )}
-          >
-            {children}
-          </div>
-          <FrontendTools tools={frontendTools} />
-        </ToolExecutionProvider>
-      </ElementsContext.Provider>
+      <ChatIdContext.Provider value={{ chatId: currentChatId }}>
+        <ElementsContext.Provider value={contextValue}>
+          <ToolExecutionProvider tools={executableTools}>
+            <div
+              className={cn(
+                ROOT_SELECTOR,
+                (contextValue?.config.variant === 'standalone' ||
+                  contextValue?.config.variant === 'sidecar') &&
+                  'h-full'
+              )}
+            >
+              {children}
+            </div>
+            <FrontendTools tools={frontendTools} />
+          </ToolExecutionProvider>
+        </ElementsContext.Provider>
+      </ChatIdContext.Provider>
     </AssistantRuntimeProvider>
   )
 }

package/src/contexts/contexts.ts

@@ -8,3 +8,5 @@ export const ElementsContext = createContext<ElementsContextType | undefined>(
 
 export const ToolApprovalContext =
   createContext<ToolApprovalContextType | null>(null)
+
+export { useChatId } from './ChatIdContext'

package/src/hooks/useAuth.ts

@@ -1,3 +1,4 @@
+import { useReplayContext } from '@/contexts/ReplayContext'
 import { hasExplicitSessionAuth, isStaticSessionAuth } from '@/lib/auth'
 import { useMemo } from 'react'
 import { ApiConfig } from '../types'
@@ -40,24 +41,38 @@ export const useAuth = ({
   auth?: ApiConfig
   projectSlug: string
 }): Auth => {
+  const replayCtx = useReplayContext()
+  const isReplay = replayCtx?.isReplay ?? false
+
   const getSession = useMemo(() => {
+    // In replay mode, skip session fetching entirely
+    if (isReplay) {
+      return null
+    }
     if (isStaticSessionAuth(auth)) {
       return () => Promise.resolve(auth.sessionToken)
     }
     return !isStaticSessionAuth(auth) && hasExplicitSessionAuth(auth)
       ? auth.sessionFn
       : defaultGetSession
-  }, [auth])
+  }, [auth, isReplay])
 
-  // The session request is only neccessary if we are not using an API key auth
+  // The session request is only neccessary if we are not using static session auth
   // configuration. If a custom session fetcher is provided, we use it,
   // otherwise we fallback to the default session fetcher
   const session = useSession({
-    // We want to check it's NOT API key auth, as the default auth scheme is session auth (if the user hasn't provided an explicit API config, we have a session auth config by default)
     getSession,
     projectSlug,
   })
 
+  // In replay mode, return immediately without waiting for session
+  if (isReplay) {
+    return {
+      headers: {},
+      isLoading: false,
+    }
+  }
+
   return !session
     ? {
         isLoading: true,

package/src/hooks/useFollowOnSuggestions.ts

@@ -1,3 +1,4 @@
+import { useReplayContext } from '@/contexts/ReplayContext'
 import { useAssistantState } from '@assistant-ui/react'
 import { useCallback, useEffect, useRef, useState } from 'react'
 import { useElements } from './useElements'
@@ -37,13 +38,17 @@ export function useFollowOnSuggestions(): {
   isLoading: boolean
 } {
   const { config } = useElements()
+  const replayCtx = useReplayContext()
+  const isReplay = replayCtx?.isReplay ?? false
+
   const auth = useAuth({
     auth: config.api,
     projectSlug: config.projectSlug,
   })
 
   // Check if follow-up suggestions are enabled (default: true)
-  const isEnabled = config.thread?.followUpSuggestions !== false
+  // Disable in replay mode since we don't need AI-generated suggestions
+  const isEnabled = !isReplay && config.thread?.followUpSuggestions !== false
 
   const [suggestions, setSuggestions] = useState<FollowOnSuggestion[]>([])
   const [isLoading, setIsLoading] = useState(false)

package/src/index.ts

@@ -10,6 +10,7 @@ export { ElementsProvider } from './contexts/ElementsProvider'
 export { useElements as useGramElements } from './hooks/useElements'
 export { useElements } from './hooks/useElements'
 export { useThreadId } from './hooks/useThreadId'
+export { useChatId } from './contexts/ChatIdContext'
 
 // Core Components
 export { Chat } from '@/components/Chat'

package/src/plugins/chart/catalog.ts

@@ -0,0 +1,141 @@
+import { createCatalog } from '@json-render/core'
+import { z } from 'zod'
+
+/**
+ * Data point schema - common structure for all chart types
+ */
+const dataPointSchema = z.object({
+  label: z.string(),
+  value: z.number(),
+  color: z.string().optional(),
+})
+
+/**
+ * Multi-series data point for line/area charts
+ */
+const seriesDataPointSchema = z
+  .object({
+    label: z.string(),
+  })
+  .catchall(z.number())
+
+/**
+ * Chart Catalog
+ *
+ * Defines all available chart components for LLM-generated visualizations.
+ * Uses Recharts under the hood for rendering.
+ */
+export const chartCatalog = createCatalog({
+  name: 'chart',
+  components: {
+    BarChart: {
+      props: z.object({
+        title: z.string().optional(),
+        data: z.array(dataPointSchema),
+        layout: z.enum(['vertical', 'horizontal']).optional(),
+        showGrid: z.boolean().optional(),
+        showLegend: z.boolean().optional(),
+        className: z.string().optional(),
+      }),
+      description:
+        'Bar chart for comparing categorical data. Use vertical for few categories, horizontal for many or long labels.',
+    },
+
+    LineChart: {
+      props: z.object({
+        title: z.string().optional(),
+        data: z.array(seriesDataPointSchema),
+        series: z.array(z.string()).optional(),
+        showGrid: z.boolean().optional(),
+        showLegend: z.boolean().optional(),
+        showDots: z.boolean().optional(),
+        curved: z.boolean().optional(),
+        className: z.string().optional(),
+      }),
+      description:
+        'Line chart for showing trends over time or continuous data. Supports multiple series.',
+    },
+
+    AreaChart: {
+      props: z.object({
+        title: z.string().optional(),
+        data: z.array(seriesDataPointSchema),
+        series: z.array(z.string()).optional(),
+        stacked: z.boolean().optional(),
+        showGrid: z.boolean().optional(),
+        showLegend: z.boolean().optional(),
+        className: z.string().optional(),
+      }),
+      description:
+        'Area chart for showing volume/magnitude over time. Use stacked for part-to-whole relationships.',
+    },
+
+    PieChart: {
+      props: z.object({
+        title: z.string().optional(),
+        data: z.array(dataPointSchema),
+        showLabels: z.boolean().optional(),
+        showLegend: z.boolean().optional(),
+        className: z.string().optional(),
+      }),
+      description:
+        'Pie chart for showing proportions of a whole. Best for 2-6 categories.',
+    },
+
+    DonutChart: {
+      props: z.object({
+        title: z.string().optional(),
+        data: z.array(dataPointSchema),
+        showLabels: z.boolean().optional(),
+        showLegend: z.boolean().optional(),
+        innerLabel: z.string().optional(),
+        innerValue: z.union([z.string(), z.number()]).optional(),
+        className: z.string().optional(),
+      }),
+      description:
+        'Donut chart (pie with center hole). Good for showing a key metric in the center.',
+    },
+
+    ScatterChart: {
+      props: z.object({
+        title: z.string().optional(),
+        data: z.array(
+          z.object({
+            x: z.number(),
+            y: z.number(),
+            label: z.string().optional(),
+            size: z.number().optional(),
+            color: z.string().optional(),
+          })
+        ),
+        xLabel: z.string().optional(),
+        yLabel: z.string().optional(),
+        showGrid: z.boolean().optional(),
+        className: z.string().optional(),
+      }),
+      description:
+        'Scatter plot for showing correlation between two variables.',
+    },
+
+    RadarChart: {
+      props: z.object({
+        title: z.string().optional(),
+        data: z.array(dataPointSchema),
+        showLegend: z.boolean().optional(),
+        className: z.string().optional(),
+      }),
+      description:
+        'Radar/spider chart for comparing multiple attributes. Best for 3-8 dimensions.',
+    },
+  },
+})
+
+export type ChartCatalogComponentProps = typeof chartCatalog extends {
+  components: infer C
+}
+  ? {
+      [K in keyof C]: C[K] extends { props: infer P }
+        ? z.infer<P extends z.ZodType ? P : never>
+        : never
+    }
+  : never

package/src/plugins/chart/component.tsx

@@ -1,149 +1,103 @@
 'use client'
 
 import { useDensity } from '@/hooks/useDensity'
-import { useRadius } from '@/hooks/useRadius'
 import { cn } from '@/lib/utils'
+import { isJsonRenderTree, type JsonRenderNode } from '@/lib/generative-ui'
 import { SyntaxHighlighterProps } from '@assistant-ui/react-markdown'
 import { AlertCircleIcon } from 'lucide-react'
-import { FC, useEffect, useMemo, useRef, useState } from 'react'
-import { parse, View, Warn } from 'vega'
-import { expressionInterpreter } from 'vega-interpreter'
+import { FC, useMemo } from 'react'
+import { MacOSWindowFrame } from '../components/MacOSWindowFrame'
+import { PluginLoadingState } from '../components/PluginLoadingState'
+
+// Import all chart components
+import {
+  BarChart,
+  LineChart,
+  AreaChart,
+  PieChart,
+  DonutChart,
+  ScatterChart,
+  RadarChart,
+} from './ui'
+
+const loadingMessages = [
+  'Rendering chart...',
+  'Visualizing data...',
+  'Building chart...',
+  'Processing data...',
+]
+
+function getRandomLoadingMessage() {
+  return loadingMessages[Math.floor(Math.random() * loadingMessages.length)]
+}
+
+/**
+ * Chart components registry
+ */
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+const chartComponents: Record<string, FC<any>> = {
+  BarChart,
+  LineChart,
+  AreaChart,
+  PieChart,
+  DonutChart,
+  ScatterChart,
+  RadarChart,
+}
+
+/**
+ * Render a chart node from json-render tree
+ */
+function renderChartNode(node: JsonRenderNode): React.ReactNode {
+  const Component = chartComponents[node.type]
+
+  if (!Component) {
+    return (
+      <div className="text-muted-foreground flex items-center gap-2 text-sm">
+        <AlertCircleIcon className="size-4" />
+        <span>Unknown chart type: {node.type}</span>
+      </div>
+    )
+  }
+
+  return <Component {...(node.props ?? {})} />
+}
 
 export const ChartRenderer: FC<SyntaxHighlighterProps> = ({ code }) => {
-  const containerRef = useRef<HTMLDivElement>(null)
-  const viewRef = useRef<View | null>(null)
-  const wrapperRef = useRef<HTMLDivElement>(null)
-  const [error, setError] = useState<string | null>(null)
-  const [chartReady, setChartReady] = useState(false)
-  const [containerWidth, setContainerWidth] = useState(0)
-  const r = useRadius()
   const d = useDensity()
 
-  // Track container width so the Vega view can fill available space
-  useEffect(() => {
-    const el = wrapperRef.current
-    if (!el) return
-    const ro = new ResizeObserver(([entry]) => {
-      setContainerWidth(entry.contentRect.width)
-    })
-    ro.observe(el)
-    return () => ro.disconnect()
-  }, [])
-
-  // Parse and validate JSON in useMemo - only recomputes when code changes
-  const parsedSpec = useMemo(() => {
+  // Parse JSON - returns null if invalid (still streaming)
+  const content = useMemo(() => {
     const trimmedCode = code.trim()
     if (!trimmedCode) return null
 
     try {
-      const spec = JSON.parse(trimmedCode) as Record<string, unknown>
-
-      // Validate that data array exists and has at least one record with values
-      const dataArray = spec.data as Array<{ values?: unknown[] }> | undefined
-      if (!dataArray?.length) return null
-
-      const hasValidData = dataArray.some(
-        (d) => Array.isArray(d.values) && d.values.length > 0
-      )
-      if (!hasValidData) return null
-
-      return spec
+      const parsed = JSON.parse(trimmedCode)
+      // Validate it has a type field (basic json-render structure)
+      if (!isJsonRenderTree(parsed)) {
+        return null
+      }
+      return parsed
     } catch {
+      // JSON is incomplete (still streaming) - return null to show loading state
       return null
     }
   }, [code])
 
-  // Only render when we have valid JSON
-  const shouldRender = parsedSpec !== null
-
-  // Build the spec with autosize and width derived from the container
-  const sizedSpec = useMemo(() => {
-    if (!parsedSpec || containerWidth === 0) return null
-    // Padding used by the outer wrapper (p-lg ≈ 24px each side)
-    const padding = 48
-    const availableWidth = Math.max(containerWidth - padding, 100)
-    return {
-      ...parsedSpec,
-      width: availableWidth,
-      autosize: { type: 'fit' as const, contains: 'padding' as const },
-    }
-  }, [parsedSpec, containerWidth])
-
-  useEffect(() => {
-    if (!containerRef.current || !shouldRender || !sizedSpec) {
-      return
-    }
-
-    setError(null)
-    setChartReady(false)
-
-    const runChart = async () => {
-      try {
-        // Clean up any existing view
-        if (viewRef.current) {
-          viewRef.current.finalize()
-          viewRef.current = null
-        }
-
-        const chart = parse(sizedSpec, undefined, { ast: true })
-        const view = new View(chart, {
-          container: containerRef.current ?? undefined,
-          renderer: 'svg',
-          hover: true,
-          logLevel: Warn,
-          expr: expressionInterpreter,
-        })
-        viewRef.current = view
-
-        await view.runAsync()
-        setChartReady(true)
-      } catch (err) {
-        console.error('Failed to render chart:', err)
-        setError(err instanceof Error ? err.message : 'Failed to render chart')
-      }
-    }
+  // Memoize the loading message so it doesn't change on every render
+  const loadingMessage = useMemo(() => getRandomLoadingMessage(), [])
 
-    runChart()
-
-    return () => {
-      if (viewRef.current) {
-        viewRef.current.finalize()
-        viewRef.current = null
-      }
-    }
-  }, [shouldRender, sizedSpec])
-
-  const showLoading = !chartReady && !error
+  // Show loading shimmer while JSON is incomplete/streaming
+  if (!content) {
+    return <PluginLoadingState text={loadingMessage} />
+  }
 
+  // Render with macOS-style window frame
   return (
-    <div
-      ref={wrapperRef}
-      className={cn(
-        // the after:hidden is to prevent assistant-ui from showing its default code block loading indicator
-        'border-border relative min-h-[400px] w-full overflow-hidden border after:hidden',
-        r('lg'),
-        showLoading ? '' : d('p-lg')
-      )}
-    >
-      {showLoading && (
-        <div className="bg-muted absolute inset-0 z-10 flex items-center justify-center">
-          <span className="shimmer text-muted-foreground text-sm">
-            Rendering chart...
-          </span>
-        </div>
-      )}
-
-      {error && (
-        <div className="bg-background absolute inset-0 z-10 flex items-center justify-center gap-2 text-rose-500">
-          <AlertCircleIcon name="alert-circle" className="h-4 w-4" />
-          {error}
-        </div>
-      )}
-
-      <div
-        ref={containerRef}
-        className={error || showLoading ? 'invisible' : 'block'}
-      />
-    </div>
+    <MacOSWindowFrame>
+      <div className={cn('bg-card w-full', d('p-lg'))}>
+        {renderChartNode(content)}
+      </div>
+    </MacOSWindowFrame>
   )
 }