diff --git a/src/components/NotebookPicker.tsx b/src/components/NotebookPicker.tsx
index a7f15f6..d683106 100644
--- a/src/components/NotebookPicker.tsx
+++ b/src/components/NotebookPicker.tsx
@@ -2,6 +2,7 @@ import React from 'react';
 import { ReactWidget } from '@jupyterlab/apputils';
 import { NotebookPanel } from '@jupyterlab/notebook';
 import { HTMLSelect } from '@jupyterlab/ui-components';
+import { deepnoteMetadataSchema } from '../types';
 
 export class NotebookPicker extends ReactWidget {
   private selected: string | null = null;
@@ -31,22 +32,30 @@ export class NotebookPicker extends ReactWidget {
 
     const selected = event.target.value;
     const deepnoteMetadata = this.panel.context.model.getMetadata('deepnote');
-    const notebooks = deepnoteMetadata?.notebooks;
+    const deepnoteMetadataValidated =
+      deepnoteMetadataSchema.safeParse(deepnoteMetadata);
 
-    if (notebooks && selected in notebooks) {
-      // clone the notebook JSON
-      const newModelData = { ...notebooks[selected] };
+    if (!deepnoteMetadataValidated.success) {
+      console.error(
+        'Invalid deepnote metadata:',
+        deepnoteMetadataValidated.error
+      );
+      return;
+    }
 
-      // preserve deepnote metadata *without* re-inserting all notebooks
-      newModelData.metadata = {
-        ...(newModelData.metadata ?? {}),
-        deepnote: {
-          notebook_names: deepnoteMetadata?.notebook_names ?? [],
-          notebooks: deepnoteMetadata?.notebooks ?? {}
-        }
-      };
+    const notebooks = deepnoteMetadataValidated.data.notebooks;
 
-      model.fromJSON(newModelData);
+    if (selected in notebooks) {
+      model.fromJSON({
+        cells: notebooks[selected]?.cells ?? [],
+        metadata: {
+          deepnote: {
+            notebooks
+          }
+        },
+        nbformat: 4,
+        nbformat_minor: 0
+      });
       model.dirty = false;
     }
 
@@ -56,12 +65,13 @@ export class NotebookPicker extends ReactWidget {
 
   render(): JSX.Element {
     const deepnoteMetadata = this.panel.context.model.getMetadata('deepnote');
-    const metadataNames = deepnoteMetadata?.notebook_names;
-    const names =
-      Array.isArray(metadataNames) &&
-      metadataNames.every(n => typeof n === 'string')
-        ? metadataNames
-        : [];
+
+    const deepnoteMetadataValidated =
+      deepnoteMetadataSchema.safeParse(deepnoteMetadata);
+
+    const names = deepnoteMetadataValidated.success
+      ? Object.values(deepnoteMetadataValidated.data.notebooks).map(n => n.name)
+      : [];
 
     return (
       <HTMLSelect
diff --git a/src/deepnote-content-provider.ts b/src/deepnote-content-provider.ts
index 86ec0e0..540abb6 100644
--- a/src/deepnote-content-provider.ts
+++ b/src/deepnote-content-provider.ts
@@ -4,7 +4,7 @@ import { transformDeepnoteYamlToNotebookContent } from './transform-deepnote-yam
 
 export const deepnoteContentProviderName = 'deepnote-content-provider';
 
-const deepnoteNotebookSchema = z.object({
+const deepnoteFileFromServerSchema = z.object({
   cells: z.array(z.any()), // or refine further with nbformat
   metadata: z.object({
     deepnote: z.object({
@@ -29,7 +29,7 @@ export class DeepnoteContentProvider extends RestContentProvider {
       return model;
     }
 
-    const validatedModelContent = deepnoteNotebookSchema.safeParse(
+    const validatedModelContent = deepnoteFileFromServerSchema.safeParse(
       model.content
     );
 
@@ -43,6 +43,7 @@ export class DeepnoteContentProvider extends RestContentProvider {
       return model;
     }
 
+    // Transform the Deepnote YAML to Jupyter notebook content
     const transformedModelContent =
       await transformDeepnoteYamlToNotebookContent(
         validatedModelContent.data.metadata.deepnote.rawYamlString
diff --git a/src/fallback-data.ts b/src/fallback-data.ts
index 7d3f45c..28e89cf 100644
--- a/src/fallback-data.ts
+++ b/src/fallback-data.ts
@@ -21,8 +21,7 @@ export const blankDeepnoteNotebookContent: IDeepnoteNotebookContent = {
   ],
   metadata: {
     deepnote: {
-      rawYamlString: null,
-      deepnoteFile: null
+      notebooks: {}
     }
   },
   nbformat: 4,
diff --git a/src/transform-deepnote-yaml-to-notebook-content.ts b/src/transform-deepnote-yaml-to-notebook-content.ts
index 83e7fdc..a760ab8 100644
--- a/src/transform-deepnote-yaml-to-notebook-content.ts
+++ b/src/transform-deepnote-yaml-to-notebook-content.ts
@@ -1,4 +1,4 @@
-import { IDeepnoteNotebookContent } from './types';
+import { IDeepnoteNotebookContent, IDeepnoteNotebookMetadata } from './types';
 import { blankCodeCell, blankDeepnoteNotebookContent } from './fallback-data';
 import { deserializeDeepnoteFile } from '@deepnote/blocks';
 import { convertDeepnoteBlockToJupyterCell } from './convert-deepnote-block-to-jupyter-cell';
@@ -9,6 +9,18 @@ export async function transformDeepnoteYamlToNotebookContent(
   try {
     const deepnoteFile = await deserializeDeepnoteFile(yamlString);
 
+    const notebooks = deepnoteFile.project.notebooks.reduce(
+      (acc, notebook) => {
+        acc[notebook.name] = {
+          id: notebook.id,
+          name: notebook.name,
+          cells: notebook.blocks.map(convertDeepnoteBlockToJupyterCell)
+        };
+        return acc;
+      },
+      {} as IDeepnoteNotebookMetadata['deepnote']['notebooks']
+    );
+
     const selectedNotebook = deepnoteFile.project.notebooks[0];
 
     if (!selectedNotebook) {
@@ -28,8 +40,14 @@ export async function transformDeepnoteYamlToNotebookContent(
     );
 
     return {
-      ...blankDeepnoteNotebookContent,
-      cells
+      cells,
+      metadata: {
+        deepnote: {
+          notebooks
+        }
+      },
+      nbformat: 4,
+      nbformat_minor: 0
     };
   } catch (error) {
     console.error('Failed to deserialize Deepnote file:', error);
diff --git a/src/types.ts b/src/types.ts
index 73f8435..95ad5e8 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -1,10 +1,33 @@
-import { INotebookContent, INotebookMetadata } from '@jupyterlab/nbformat';
-import type { DeepnoteFile } from '@deepnote/blocks';
+import {
+  ICodeCell,
+  IMarkdownCell,
+  INotebookContent,
+  INotebookMetadata
+} from '@jupyterlab/nbformat';
+
+import { z } from 'zod';
+
+export const deepnoteMetadataSchema = z.object({
+  notebooks: z.record(
+    z.string(),
+    z.object({
+      id: z.string(),
+      name: z.string(),
+      cells: z.array(z.any())
+    })
+  )
+});
 
 export interface IDeepnoteNotebookMetadata extends INotebookMetadata {
   deepnote: {
-    rawYamlString: string | null;
-    deepnoteFile: DeepnoteFile | null;
+    notebooks: Record<
+      string,
+      {
+        id: string;
+        name: string;
+        cells: Array<ICodeCell | IMarkdownCell>;
+      }
+    >;
   };
 }