feat: better heatmap styling

src/main/a2ui/catalog.ts

@@ -15,7 +15,7 @@ const CATALOG_ENTRIES: A2UICatalogEntry[] = [
   { type: 'text', description: 'Text block with Markdown support' },
   { type: 'button', description: 'Clickable button that dispatches an action' },
   { type: 'card', description: 'Card with title, subtitle, body, and action buttons' },
-  { type: 'chart', description: 'Bar, line, or pie chart visualization', custom: true },
+  { type: 'chart', description: 'Bar, stacked-bar, line, area, pie, donut, or heatmap chart visualization', custom: true },
   { type: 'table', description: 'Data table with columns and rows', custom: true },
   { type: 'textField', description: 'Text input field with data binding' },
   { type: 'checkBox', description: 'Checkbox input with data binding' },

src/main/a2ui/generator.ts

@@ -56,7 +56,7 @@ function createSurfaceMessages(
 // ---- Tool argument interfaces ----
 
 export interface RenderChartArgs {
-  chartType: 'bar' | 'stacked-bar' | 'line' | 'pie';
+  chartType: 'bar' | 'stacked-bar' | 'line' | 'area' | 'pie' | 'donut' | 'heatmap';
   title?: string;
   series: Array<{ label: string; value: number; segments?: Array<{ label: string; value: number }> }>;
 }

src/main/engine/ChatEngine.ts

@@ -323,7 +323,7 @@ Available Data Tools:
 - get_media_posts: Get posts that use a specific media file.
 
 Available UI Render Tools (use these to show rich interactive elements):
-- render_chart: Show data as a bar, stacked-bar, line, or pie chart. Use when presenting statistics or comparisons. Use stacked-bar when each bar has multiple segments (e.g., published vs draft posts per year).
+- render_chart: Show data as a bar, stacked-bar, line, area, pie, donut, or heatmap chart. Use when presenting statistics or comparisons. Use stacked-bar when each bar has multiple segments (e.g., published vs draft posts per year). Use area for cumulative or trend data where the filled region emphasizes volume. Use donut for proportional breakdowns with a total displayed in the center. Use heatmap for grid/matrix visualizations where color intensity shows magnitude — e.g., posts per month across years (each series entry is a row like a year, each segment is a column like a month), or a calendar view where rows are weekdays and columns are week numbers. ALWAYS prefer heatmap over a table with emojis or color indicators when showing intensity grids or calendar-style activity views.
 - render_table: Show data in a structured table. Use for tabular comparisons and listings.
 - render_form: Show an interactive form to collect user input (e.g., metadata edits, settings).
 - render_card: Show an information card with title, body, and action buttons.

src/main/engine/OpenCodeManager.ts

@@ -924,31 +924,31 @@ export class OpenCodeManager {
         input_schema: {
           type: 'object',
           properties: {
-            chartType: { type: 'string', enum: ['bar', 'stacked-bar', 'line', 'pie'], description: 'The type of chart to render. Use stacked-bar when each bar has multiple segments (categories).' },
+            chartType: { type: 'string', enum: ['bar', 'stacked-bar', 'line', 'area', 'pie', 'donut', 'heatmap'], description: 'The type of chart to render. Use stacked-bar when each bar has multiple segments (categories). Use area for trend/cumulative data. Use donut for proportional data with a total in the center. Use heatmap for grid/matrix visualizations where color intensity shows magnitude (e.g., posts per month across years). Prefer heatmap over tables with emojis for intensity data.' },
             title: { type: 'string', description: 'Optional chart title' },
             series: {
               type: 'array',
               items: {
                 type: 'object',
                 properties: {
-                  label: { type: 'string', description: 'Data point label' },
-                  value: { type: 'number', description: 'Data point value (total for stacked bars)' },
+                  label: { type: 'string', description: 'Data point label (row label for heatmaps, e.g., year)' },
+                  value: { type: 'number', description: 'Data point value (total for stacked bars, ignored for heatmaps)' },
                   segments: {
                     type: 'array',
                     items: {
                       type: 'object',
                       properties: {
-                        label: { type: 'string', description: 'Segment category label' },
-                        value: { type: 'number', description: 'Segment value' },
+                        label: { type: 'string', description: 'Segment/column label (e.g., month name for heatmaps)' },
+                        value: { type: 'number', description: 'Segment value (color intensity for heatmaps)' },
                       },
                       required: ['label', 'value'],
                     },
-                    description: 'Segments for stacked-bar charts. Each segment is a category within the bar.',
+                    description: 'Segments within this data point. Required for stacked-bar and heatmap charts. For heatmaps, each segment becomes a cell in that row.',
                   },
                 },
                 required: ['label', 'value'],
               },
-              description: 'Array of data points with label and value. For stacked-bar charts, include segments.',
+              description: 'Array of data points. For stacked-bar and heatmap charts, include segments. For heatmaps, each entry is a row and segments are columns.',
             },
           },
           required: ['chartType', 'series'],
@@ -1070,12 +1070,24 @@ export class OpenCodeManager {
                         value: { type: 'string', description: 'Display value (for type metric)' },
                         title: { type: 'string', description: 'Title (for type list, chart, or table)' },
                         items: { type: 'array', items: { type: 'string' }, description: 'Items (for type list)' },
-                        chartType: { type: 'string', enum: ['bar', 'line', 'pie'], description: 'Chart type (for type chart)' },
+                        chartType: { type: 'string', enum: ['bar', 'stacked-bar', 'line', 'area', 'pie', 'donut', 'heatmap'], description: 'Chart type (for type chart). Use heatmap for intensity grids.' },
                         series: {
                           type: 'array',
                           items: {
                             type: 'object',
-                            properties: { label: { type: 'string' }, value: { type: 'number' } },
+                            properties: {
+                              label: { type: 'string' },
+                              value: { type: 'number' },
+                              segments: {
+                                type: 'array',
+                                items: {
+                                  type: 'object',
+                                  properties: { label: { type: 'string' }, value: { type: 'number' } },
+                                  required: ['label', 'value'],
+                                },
+                                description: 'Segments for stacked-bar and heatmap charts',
+                              },
+                            },
                             required: ['label', 'value'],
                           },
                           description: 'Data series (for type chart)',

src/renderer/a2ui/components/A2UIChart.tsx

@@ -33,6 +33,76 @@ function getSegmentColor(index: number): string {
   return SEGMENT_COLORS[index % SEGMENT_COLORS.length];
 }
 
+/* ── Heatmap colour helpers ── */
+
+type RGB = [number, number, number];
+
+const FALLBACK_INS: RGB = [53, 117, 56];
+const FALLBACK_DEL: RGB = [183, 72, 72];
+
+/** Parse a CSS color value (hex or rgb/rgba) into an [r,g,b] triple. */
+function parseCssColor(raw: string): RGB | null {
+  const v = raw.trim();
+  const hexMatch = v.match(/^#([0-9a-f]{6})$/i);
+  if (hexMatch) {
+    const hex = hexMatch[1];
+    return [
+      Number.parseInt(hex.slice(0, 2), 16),
+      Number.parseInt(hex.slice(2, 4), 16),
+      Number.parseInt(hex.slice(4, 6), 16),
+    ];
+  }
+  const rgbMatch = v.match(/^rgba?\(([^)]+)\)$/i);
+  if (rgbMatch) {
+    const channels = rgbMatch[1].split(',').map((c) => Math.round(Number.parseFloat(c.trim()))).slice(0, 3);
+    if (channels.length === 3 && channels.every((c) => Number.isFinite(c))) {
+      return channels.map((c) => Math.max(0, Math.min(255, c))) as RGB;
+    }
+  }
+  return null;
+}
+
+function lerpRGB(a: RGB, b: RGB, t: number): RGB {
+  return [
+    Math.round(a[0] + (b[0] - a[0]) * t),
+    Math.round(a[1] + (b[1] - a[1]) * t),
+    Math.round(a[2] + (b[2] - a[2]) * t),
+  ];
+}
+
+/** Relative luminance (0-255 scale) for contrast decision. */
+function luminance(c: RGB): number {
+  return 0.299 * c[0] + 0.587 * c[1] + 0.114 * c[2];
+}
+
+function readPicoInsDelColors(): { ins: RGB; del: RGB } {
+  try {
+    const style = window.getComputedStyle(document.documentElement);
+    const ins = parseCssColor(style.getPropertyValue('--pico-ins-color')) ?? FALLBACK_INS;
+    const del = parseCssColor(style.getPropertyValue('--pico-del-color')) ?? FALLBACK_DEL;
+    return { ins, del };
+  } catch {
+    return { ins: FALLBACK_INS, del: FALLBACK_DEL };
+  }
+}
+
+/** Compute heatmap cell background and contrasting text color. */
+function heatmapCellColors(alpha: number, ins: RGB, del: RGB): { bg: string; fg: string } {
+  if (alpha <= 0) return { bg: 'transparent', fg: 'inherit' };
+  const rgb = lerpRGB(ins, del, alpha);
+  // Scale opacity so even low values are visible (0.25 → 1.0)
+  const opacity = 0.25 + alpha * 0.75;
+  const bg = `rgba(${rgb[0]},${rgb[1]},${rgb[2]},${opacity.toFixed(2)})`;
+  // Blend the effective RGB with the assumed dark background (~30) for contrast calc
+  const effective: RGB = [
+    Math.round(rgb[0] * opacity + 30 * (1 - opacity)),
+    Math.round(rgb[1] * opacity + 30 * (1 - opacity)),
+    Math.round(rgb[2] * opacity + 30 * (1 - opacity)),
+  ];
+  const fg = luminance(effective) > 140 ? '#000' : '#fff';
+  return { bg, fg };
+}
+
 /** Collect unique segment labels across all series entries, preserving order. */
 function collectSegmentLabels(series: SeriesEntry[]): string[] {
   const seen = new Set<string>();
@@ -50,6 +120,297 @@ function collectSegmentLabels(series: SeriesEntry[]): string[] {
   return labels;
 }
 
+/** Generate nice round tick values for the Y-axis. */
+function computeYTicks(maxValue: number, tickCount: number = 4): number[] {
+  if (maxValue <= 0) return [0];
+  const rawStep = maxValue / (tickCount - 1);
+  const magnitude = Math.pow(10, Math.floor(Math.log10(rawStep)));
+  const normalised = rawStep / magnitude;
+  let niceStep: number;
+  if (normalised <= 1) niceStep = magnitude;
+  else if (normalised <= 2) niceStep = 2 * magnitude;
+  else if (normalised <= 5) niceStep = 5 * magnitude;
+  else niceStep = 10 * magnitude;
+
+  const ticks: number[] = [];
+  for (let v = 0; v <= maxValue + niceStep * 0.01; v += niceStep) {
+    ticks.push(Math.round(v * 1e6) / 1e6);
+  }
+  if (ticks.length < 2) ticks.push(niceStep);
+  return ticks;
+}
+
+const LINE_CHART_PADDING = { top: 8, right: 12, bottom: 24, left: 40 };
+const LINE_CHART_HEIGHT = 140;
+
+function renderLineChart(component: A2UIResolvedComponent, series: SeriesEntry[], showArea: boolean = false): React.ReactNode {
+  if (series.length === 0) return null;
+
+  const maxValue = Math.max(...series.map((e) => e.value), 0);
+  const yTicks = computeYTicks(maxValue);
+  const yMax = yTicks[yTicks.length - 1];
+
+  const pad = LINE_CHART_PADDING;
+  const plotWidth = 300 - pad.left - pad.right;
+  const plotHeight = LINE_CHART_HEIGHT - pad.top - pad.bottom;
+
+  const xStep = series.length > 1 ? plotWidth / (series.length - 1) : 0;
+
+  const points = series.map((entry, i) => {
+    const x = pad.left + (series.length > 1 ? i * xStep : plotWidth / 2);
+    const y = pad.top + (yMax > 0 ? (1 - entry.value / yMax) * plotHeight : plotHeight);
+    return { x, y, entry };
+  });
+
+  const polylinePoints = points.map((p) => `${p.x},${p.y}`).join(' ');
+  const baselineY = pad.top + plotHeight;
+  const areaPoints = showArea
+    ? `${points[0].x},${baselineY} ${polylinePoints} ${points[points.length - 1].x},${baselineY}`
+    : '';
+
+  return (
+    <svg
+      className="assistant-panel-chart-line-svg"
+      viewBox={`0 0 300 ${LINE_CHART_HEIGHT}`}
+      preserveAspectRatio="xMidYMid meet"
+    >
+      {/* Horizontal grid lines + Y-axis labels */}
+      {yTicks.map((tick) => {
+        const y = pad.top + (yMax > 0 ? (1 - tick / yMax) * plotHeight : plotHeight);
+        return (
+          <g key={`grid-${tick}`}>
+            <line
+              className="assistant-panel-chart-line-grid"
+              x1={pad.left}
+              y1={y}
+              x2={pad.left + plotWidth}
+              y2={y}
+            />
+            <text
+              className="assistant-panel-chart-line-y-label"
+              x={pad.left - 4}
+              y={y}
+              textAnchor="end"
+              dominantBaseline="middle"
+            >
+              {tick}
+            </text>
+          </g>
+        );
+      })}
+
+      {/* Area fill (when showArea is true) */}
+      {showArea && (
+        <polygon
+          className="assistant-panel-chart-area-fill"
+          points={areaPoints}
+        />
+      )}
+
+      {/* Line */}
+      <polyline
+        className="assistant-panel-chart-line-path"
+        points={polylinePoints}
+        fill="none"
+      />
+
+      {/* Dots */}
+      {points.map((p, i) => (
+        <circle
+          key={`${component.id}-dot-${i}`}
+          className="assistant-panel-chart-line-dot"
+          cx={p.x}
+          cy={p.y}
+          r={3}
+        >
+          <title>{`${p.entry.label}: ${p.entry.value}`}</title>
+        </circle>
+      ))}
+
+      {/* X-axis labels */}
+      {points.map((p, i) => (
+        <text
+          key={`${component.id}-xlabel-${i}`}
+          className="assistant-panel-chart-line-x-label"
+          x={p.x}
+          y={LINE_CHART_HEIGHT - 4}
+          textAnchor="middle"
+        >
+          {p.entry.label}
+        </text>
+      ))}
+    </svg>
+  );
+}
+
+const PIE_CHART_SIZE = 140;
+const PIE_CHART_RADIUS = 56;
+const PIE_CHART_CENTER = PIE_CHART_SIZE / 2;
+
+function describePieSlice(cx: number, cy: number, r: number, startAngle: number, endAngle: number): string {
+  const startRad = (startAngle - 90) * (Math.PI / 180);
+  const endRad = (endAngle - 90) * (Math.PI / 180);
+  const x1 = cx + r * Math.cos(startRad);
+  const y1 = cy + r * Math.sin(startRad);
+  const x2 = cx + r * Math.cos(endRad);
+  const y2 = cy + r * Math.sin(endRad);
+  const largeArc = endAngle - startAngle > 180 ? 1 : 0;
+  return `M${cx},${cy} L${x1},${y1} A${r},${r} 0 ${largeArc} 1 ${x2},${y2} Z`;
+}
+
+function renderPieChart(component: A2UIResolvedComponent, series: SeriesEntry[], isDonut: boolean = false): React.ReactNode {
+  if (series.length === 0) return null;
+
+  const total = series.reduce((sum, e) => sum + e.value, 0);
+  if (total <= 0) return null;
+
+  let currentAngle = 0;
+  const slices = series.map((entry, i) => {
+    const sliceAngle = (entry.value / total) * 360;
+    const startAngle = currentAngle;
+    const endAngle = currentAngle + sliceAngle;
+    currentAngle = endAngle;
+
+    // For a full circle (single slice), use a circle element instead
+    if (sliceAngle >= 359.99) {
+      return (
+        <circle
+          key={`${component.id}-slice-${i}`}
+          className="assistant-panel-chart-pie-slice"
+          cx={PIE_CHART_CENTER}
+          cy={PIE_CHART_CENTER}
+          r={PIE_CHART_RADIUS}
+          fill={getSegmentColor(i)}
+        >
+          <title>{`${entry.label}: ${entry.value}`}</title>
+        </circle>
+      );
+    }
+
+    return (
+      <path
+        key={`${component.id}-slice-${i}`}
+        className="assistant-panel-chart-pie-slice"
+        d={describePieSlice(PIE_CHART_CENTER, PIE_CHART_CENTER, PIE_CHART_RADIUS, startAngle, endAngle)}
+        fill={getSegmentColor(i)}
+      >
+        <title>{`${entry.label}: ${entry.value}`}</title>
+      </path>
+    );
+  });
+
+  const donutInnerRadius = PIE_CHART_RADIUS * 0.58;
+
+  return (
+    <>
+      <svg
+        className="assistant-panel-chart-pie-svg"
+        viewBox={`0 0 ${PIE_CHART_SIZE} ${PIE_CHART_SIZE}`}
+        preserveAspectRatio="xMidYMid meet"
+      >
+        {slices}
+        {isDonut && (
+          <>
+            <circle
+              className="assistant-panel-chart-donut-hole"
+              cx={PIE_CHART_CENTER}
+              cy={PIE_CHART_CENTER}
+              r={donutInnerRadius}
+            />
+            <text
+              className="assistant-panel-chart-donut-total"
+              x={PIE_CHART_CENTER}
+              y={PIE_CHART_CENTER}
+              textAnchor="middle"
+              dominantBaseline="central"
+            >
+              {total}
+            </text>
+          </>
+        )}
+      </svg>
+      <div className="assistant-panel-chart-legend">
+        {series.map((entry, i) => (
+          <span key={entry.label} className="assistant-panel-chart-legend-item">
+            <span
+              className="assistant-panel-chart-legend-swatch"
+              style={{ backgroundColor: getSegmentColor(i) }}
+            />
+            {entry.label}
+          </span>
+        ))}
+      </div>
+    </>
+  );
+}
+
+function renderHeatmap(component: A2UIResolvedComponent, series: SeriesEntry[]): React.ReactNode {
+  // Heatmap requires segments — each entry is a row, each segment is a column
+  const rows = series.filter((e) => e.segments && e.segments.length > 0);
+  if (rows.length === 0) return null;
+
+  const columnLabels = collectSegmentLabels(rows);
+  if (columnLabels.length === 0) return null;
+
+  // Find global max for normalisation
+  let globalMax = 0;
+  for (const row of rows) {
+    for (const seg of row.segments!) {
+      if (seg.value > globalMax) globalMax = seg.value;
+    }
+  }
+
+  const { ins, del } = readPicoInsDelColors();
+
+  // Build a lookup for each row's segment values by column label
+  const columnCount = columnLabels.length;
+
+  return (
+    <div
+      className="assistant-panel-chart-heatmap"
+      style={{
+        gridTemplateColumns: `auto repeat(${columnCount}, 1fr)`,
+      }}
+    >
+      {/* Header row: empty corner + column labels */}
+      <span className="assistant-panel-chart-heatmap-corner" />
+      {columnLabels.map((col) => (
+        <span key={`col-${col}`} className="assistant-panel-chart-heatmap-col-label">
+          {col}
+        </span>
+      ))}
+
+      {/* Data rows */}
+      {rows.map((row) => {
+        const segMap = new Map<string, number>();
+        for (const seg of row.segments!) {
+          segMap.set(seg.label, seg.value);
+        }
+        return (
+          <React.Fragment key={`${component.id}-row-${row.label}`}>
+            <span className="assistant-panel-chart-heatmap-row-label">{row.label}</span>
+            {columnLabels.map((col) => {
+              const val = segMap.get(col) ?? 0;
+              const alpha = globalMax > 0 ? val / globalMax : 0;
+              const { bg, fg } = heatmapCellColors(alpha, ins, del);
+              return (
+                <span
+                  key={`${component.id}-cell-${row.label}-${col}`}
+                  className="assistant-panel-chart-heatmap-cell"
+                  style={{ background: bg, color: fg }}
+                  title={String(val)}
+                >
+                  {val > 0 ? val : ''}
+                </span>
+              );
+            })}
+          </React.Fragment>
+        );
+      })}
+    </div>
+  );
+}
+
 export const A2UIChart: React.FC<A2UIComponentProps> = ({ component }) => {
   const chartType = String(component.properties.chartType ?? 'bar');
   const title = component.properties.title as string | undefined;
@@ -68,11 +429,24 @@ export const A2UIChart: React.FC<A2UIComponentProps> = ({ component }) => {
 
   const segmentLabels = isStacked ? collectSegmentLabels(series) : [];
 
+  const isLine = chartType === 'line';
+  const isArea = chartType === 'area';
+  const isPie = chartType === 'pie';
+  const isDonut = chartType === 'donut';
+  const isHeatmap = chartType === 'heatmap';
+
   return (
     <div className="assistant-panel-chart">
       {title && <p className="assistant-panel-chart-title">{title}</p>}
       <div className="assistant-panel-chart-type">{chartType}</div>
-      <div className="assistant-panel-chart-body">
+      {isPie || isDonut ? (
+        renderPieChart(component, series, isDonut)
+      ) : isHeatmap ? (
+        renderHeatmap(component, series)
+      ) : isLine || isArea ? (
+        renderLineChart(component, series, isArea)
+      ) : (
+        <div className="assistant-panel-chart-body">
         {series.map((entry, index) => {
           const totalValue = isStacked && entry.segments
             ? entry.segments.reduce((sum, s) => sum + s.value, 0)
@@ -109,6 +483,7 @@ export const A2UIChart: React.FC<A2UIComponentProps> = ({ component }) => {
           );
         })}
       </div>
+      )}
       {isStacked && segmentLabels.length > 0 && (
         <div className="assistant-panel-chart-legend">
           {segmentLabels.map((label, i) => (

src/renderer/components/AssistantSidebar/AssistantSidebar.css

@@ -240,6 +240,115 @@
   flex-shrink: 0;
 }
 
+/* Line chart */
+.assistant-panel-chart-line-svg {
+  width: 100%;
+  height: auto;
+  max-height: 180px;
+}
+
+.assistant-panel-chart-line-grid {
+  stroke: var(--vscode-panel-border, #444);
+  stroke-width: 0.5;
+  stroke-dasharray: 3 3;
+}
+
+.assistant-panel-chart-line-y-label {
+  font-size: 9px;
+  fill: var(--vscode-foreground, #ccc);
+  opacity: 0.7;
+}
+
+.assistant-panel-chart-line-path {
+  stroke: var(--vscode-charts-blue, #75beff);
+  stroke-width: 2;
+  stroke-linejoin: round;
+  stroke-linecap: round;
+}
+
+.assistant-panel-chart-line-dot {
+  fill: var(--vscode-charts-blue, #75beff);
+  stroke: var(--vscode-editor-background, #1e1e1e);
+  stroke-width: 1.5;
+}
+
+.assistant-panel-chart-line-x-label {
+  font-size: 9px;
+  fill: var(--vscode-foreground, #ccc);
+  opacity: 0.7;
+}
+
+.assistant-panel-chart-area-fill {
+  fill: var(--vscode-charts-blue, #75beff);
+  opacity: 0.15;
+}
+
+/* Pie chart */
+.assistant-panel-chart-pie-svg {
+  width: 100%;
+  max-width: 160px;
+  height: auto;
+  align-self: center;
+}
+
+.assistant-panel-chart-pie-slice {
+  stroke: var(--vscode-editor-background, #1e1e1e);
+  stroke-width: 1.5;
+}
+
+/* Donut chart */
+.assistant-panel-chart-donut-hole {
+  fill: var(--vscode-editor-background, #1e1e1e);
+}
+
+.assistant-panel-chart-donut-total {
+  font-size: 16px;
+  font-weight: 600;
+  fill: var(--vscode-foreground, #ccc);
+  font-variant-numeric: tabular-nums;
+}
+
+/* Heatmap chart */
+.assistant-panel-chart-heatmap {
+  display: grid;
+  gap: 2px;
+  font-size: 11px;
+}
+
+.assistant-panel-chart-heatmap-corner {
+  /* empty top-left cell */
+}
+
+.assistant-panel-chart-heatmap-col-label {
+  text-align: center;
+  opacity: 0.7;
+  white-space: nowrap;
+  overflow: hidden;
+  text-overflow: ellipsis;
+}
+
+.assistant-panel-chart-heatmap-row-label {
+  text-align: right;
+  padding-right: 4px;
+  opacity: 0.7;
+  white-space: nowrap;
+  overflow: hidden;
+  text-overflow: ellipsis;
+  max-width: 80px;
+}
+
+.assistant-panel-chart-heatmap-cell {
+  aspect-ratio: 1;
+  min-width: 14px;
+  min-height: 14px;
+  border-radius: 2px;
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  font-size: 10px;
+  font-weight: 500;
+}
+
 .assistant-panel-form {
   display: flex;
   flex-direction: column;

src/renderer/navigation/assistantPanelSpec.ts

@@ -30,14 +30,20 @@ const actionElementSchema = z.object({
   payload: z.record(z.string(), z.unknown()).optional(),
 });
 
+const segmentSchema = z.object({
+  label: z.string().min(1),
+  value: z.number(),
+});
+
 const chartElementSchema = z.object({
   type: z.literal('chart'),
-  chartType: z.enum(['bar', 'line', 'pie']),
+  chartType: z.enum(['bar', 'stacked-bar', 'line', 'area', 'pie', 'donut', 'heatmap']),
   title: z.string().min(1).optional(),
   series: z.array(
     z.object({
       label: z.string().min(1),
       value: z.number(),
+      segments: z.array(segmentSchema).optional(),
     }),
   ).min(1),
 });