diff --git a/packages/app/src/components/DBEditTimeChartForm/ChartEditorControls.tsx b/packages/app/src/components/DBEditTimeChartForm/ChartEditorControls.tsx
index acb2a9b7..9cc73983 100644
--- a/packages/app/src/components/DBEditTimeChartForm/ChartEditorControls.tsx
+++ b/packages/app/src/components/DBEditTimeChartForm/ChartEditorControls.tsx
@@ -99,6 +99,11 @@ export function ChartEditorControls({
             control={control}
             name="source"
             data-testid="source-selector"
+            allowedSourceKinds={
+              displayType === DisplayType.Heatmap
+                ? [SourceKind.Trace]
+                : undefined
+            }
             sourceSchemaPreview={
               <SourceSchemaPreview source={tableSource} variant="text" />
             }
diff --git a/packages/app/src/components/DBEditTimeChartForm/EditTimeChartForm.tsx b/packages/app/src/components/DBEditTimeChartForm/EditTimeChartForm.tsx
index 1bed1fef..6b35a513 100644
--- a/packages/app/src/components/DBEditTimeChartForm/EditTimeChartForm.tsx
+++ b/packages/app/src/components/DBEditTimeChartForm/EditTimeChartForm.tsx
@@ -417,14 +417,11 @@ export default function EditTimeChartForm({
       }
 
       if (displayType === DisplayType.Heatmap && Array.isArray(select)) {
-        const traceSource =
+        const defaultValue =
           tableSource?.kind === SourceKind.Trace &&
           tableSource.durationExpression
-            ? tableSource
-            : undefined;
-        const defaultValue = traceSource
-          ? getDurationMsExpression(traceSource)
-          : (select[0]?.valueExpression ?? '');
+            ? getDurationMsExpression(tableSource)
+            : (select[0]?.valueExpression ?? '');
         const heatmapSeries: SavedChartConfigWithSelectArray['select'] = [
           {
             aggFn: 'count',
@@ -436,12 +433,10 @@ export default function EditTimeChartForm({
         setValue('select', heatmapSeries);
         setValue('series', heatmapSeries);
         setValue('series.0.countExpression' as any, 'count()');
-        if (traceSource) {
-          setValue('numberFormat', {
-            output: 'duration' as any,
-            factor: 0.001,
-          });
-        }
+        setValue('numberFormat', {
+          output: 'duration' as any,
+          factor: 0.001,
+        });
       }
 
       // Don't auto-submit when config type changes, to avoid clearing form state (like source)