diff --git a/extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx b/extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx
index 418b5e2e4..87b070273 100644
--- a/extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx
+++ b/extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx
@@ -19,6 +19,13 @@ import { percentFormatter } from "./formatters";
 import { Mode } from "../../data-extensions-editor/shared/mode";
 import { groupMethods } from "../../data-extensions-editor/shared/sorting";
 
+const LoadingContainer = styled.div`
+  text-align: center;
+  padding: 1em;
+  font-size: x-large;
+  font-weight: 600;
+`;
+
 const DataExtensionsEditorContainer = styled.div`
   margin-top: 1rem;
 `;
@@ -239,6 +246,10 @@ export function DataExtensionsEditor({
     });
   }, [viewState?.mode]);
 
+  if (viewState === undefined) {
+    return <LoadingContainer>Loading...</LoadingContainer>;
+  }
+
   return (
     <DataExtensionsEditorContainer>
       {progress.maxStep > 0 && (
@@ -252,25 +263,21 @@ export function DataExtensionsEditor({
         <>
           <ViewTitle>Data extensions editor</ViewTitle>
           <DetailsContainer>
-            {viewState?.extensionPack && (
-              <>
-                <LinkIconButton onClick={onOpenExtensionPackClick}>
-                  <span slot="start" className="codicon codicon-package"></span>
-                  {viewState.extensionPack.name}
-                </LinkIconButton>
-              </>
-            )}
+            <LinkIconButton onClick={onOpenExtensionPackClick}>
+              <span slot="start" className="codicon codicon-package"></span>
+              {viewState.extensionPack.name}
+            </LinkIconButton>
             <div>
               {percentFormatter.format(modeledPercentage / 100)} modeled
             </div>
             <div>
               {percentFormatter.format(unModeledPercentage / 100)} unmodeled
             </div>
-            {viewState?.enableFrameworkMode && (
+            {viewState.enableFrameworkMode && (
               <>
                 <div>
                   Mode:{" "}
-                  {viewState?.mode === Mode.Framework
+                  {viewState.mode === Mode.Framework
                     ? "Framework"
                     : "Application"}
                 </div>
@@ -290,17 +297,17 @@ export function DataExtensionsEditor({
           <EditorContainer>
             <ButtonsContainer>
               <VSCodeButton onClick={onSaveAllClick}>Apply</VSCodeButton>
-              {viewState?.enableFrameworkMode && (
+              {viewState.enableFrameworkMode && (
                 <VSCodeButton appearance="secondary" onClick={onRefreshClick}>
                   Refresh
                 </VSCodeButton>
               )}
               <VSCodeButton onClick={onGenerateFromSourceClick}>
-                {viewState?.mode === Mode.Framework
+                {viewState.mode === Mode.Framework
                   ? "Generate"
                   : "Download and generate"}
               </VSCodeButton>
-              {viewState?.showLlmButton && (
+              {viewState.showLlmButton && (
                 <>
                   <VSCodeButton onClick={onGenerateAllFromLlmClick}>
                     Generate using LLM
@@ -313,7 +320,6 @@ export function DataExtensionsEditor({
               unsavedModels={unsavedModels}
               modeledMethods={modeledMethods}
               viewState={viewState}
-              mode={viewState?.mode ?? Mode.Application}
               onChange={onChange}
               onSaveModelClick={onSaveModelClick}
               onGenerateFromLlmClick={onGenerateFromLlmClick}
diff --git a/extensions/ql-vscode/src/view/data-extensions-editor/LibraryRow.tsx b/extensions/ql-vscode/src/view/data-extensions-editor/LibraryRow.tsx
index 58d8f16cf..89e50a30c 100644
--- a/extensions/ql-vscode/src/view/data-extensions-editor/LibraryRow.tsx
+++ b/extensions/ql-vscode/src/view/data-extensions-editor/LibraryRow.tsx
@@ -70,8 +70,7 @@ type Props = {
   title: string;
   externalApiUsages: ExternalApiUsage[];
   modeledMethods: Record<string, ModeledMethod>;
-  viewState: DataExtensionEditorViewState | undefined;
-  mode: Mode;
+  viewState: DataExtensionEditorViewState;
   hasUnsavedChanges: boolean;
   onChange: (
     modelName: string,
@@ -95,7 +94,6 @@ export const LibraryRow = ({
   externalApiUsages,
   modeledMethods,
   viewState,
-  mode,
   hasUnsavedChanges,
   onChange,
   onSaveModelClick,
@@ -166,7 +164,7 @@ export const LibraryRow = ({
           </ModeledPercentage>
           {hasUnsavedChanges ? <VSCodeTag>UNSAVED</VSCodeTag> : null}
         </NameContainer>
-        {viewState?.showLlmButton && (
+        {viewState.showLlmButton && (
           <VSCodeButton appearance="icon" onClick={handleModelWithAI}>
             <Codicon name="lightbulb-autofix" label="Model with AI" />
             &nbsp;Model with AI
@@ -176,8 +174,8 @@ export const LibraryRow = ({
           <Codicon name="code" label="Model from source" />
           &nbsp;Model from source
         </VSCodeButton>
-        {viewState?.enableFrameworkMode &&
-          viewState?.mode === Mode.Application && (
+        {viewState.enableFrameworkMode &&
+          viewState.mode === Mode.Application && (
             <VSCodeButton appearance="icon" onClick={handleModelDependency}>
               <Codicon name="references" label="Model dependency" />
               &nbsp;Model dependency
@@ -190,7 +188,7 @@ export const LibraryRow = ({
           <ModeledMethodDataGrid
             externalApiUsages={externalApiUsages}
             modeledMethods={modeledMethods}
-            mode={mode}
+            mode={viewState.mode}
             onChange={onChangeWithModelName}
           />
           <SectionDivider />
diff --git a/extensions/ql-vscode/src/view/data-extensions-editor/ModeledMethodsList.tsx b/extensions/ql-vscode/src/view/data-extensions-editor/ModeledMethodsList.tsx
index 18ff58709..a70cbc84f 100644
--- a/extensions/ql-vscode/src/view/data-extensions-editor/ModeledMethodsList.tsx
+++ b/extensions/ql-vscode/src/view/data-extensions-editor/ModeledMethodsList.tsx
@@ -3,7 +3,6 @@ import { useMemo } from "react";
 import { ExternalApiUsage } from "../../data-extensions-editor/external-api-usage";
 import { ModeledMethod } from "../../data-extensions-editor/modeled-method";
 import { LibraryRow } from "./LibraryRow";
-import { Mode } from "../../data-extensions-editor/shared/mode";
 import {
   groupMethods,
   sortGroupNames,
@@ -14,8 +13,7 @@ type Props = {
   externalApiUsages: ExternalApiUsage[];
   unsavedModels: Set<string>;
   modeledMethods: Record<string, ModeledMethod>;
-  viewState: DataExtensionEditorViewState | undefined;
-  mode: Mode;
+  viewState: DataExtensionEditorViewState;
   onChange: (
     modelName: string,
     externalApiUsage: ExternalApiUsage,
@@ -38,15 +36,14 @@ export const ModeledMethodsList = ({
   unsavedModels,
   modeledMethods,
   viewState,
-  mode,
   onChange,
   onSaveModelClick,
   onGenerateFromLlmClick,
   onGenerateFromSourceClick,
 }: Props) => {
   const grouped = useMemo(
-    () => groupMethods(externalApiUsages, mode),
-    [externalApiUsages, mode],
+    () => groupMethods(externalApiUsages, viewState.mode),
+    [externalApiUsages, viewState.mode],
   );
 
   const sortedGroupNames = useMemo(() => sortGroupNames(grouped), [grouped]);
@@ -61,7 +58,6 @@ export const ModeledMethodsList = ({
           hasUnsavedChanges={unsavedModels.has(libraryName)}
           modeledMethods={modeledMethods}
           viewState={viewState}
-          mode={mode}
           onChange={onChange}
           onSaveModelClick={onSaveModelClick}
           onGenerateFromLlmClick={onGenerateFromLlmClick}