diff --git a/src/vs/workbench/api/browser/mainThreadNotebook.ts b/src/vs/workbench/api/browser/mainThreadNotebook.ts
index 8163fe1bde0..e195b98f0dc 100644
--- a/src/vs/workbench/api/browser/mainThreadNotebook.ts
+++ b/src/vs/workbench/api/browser/mainThreadNotebook.ts
@@ -152,6 +152,10 @@ export class MainThreadNotebooks extends Disposable implements MainThreadNoteboo
 	}
 
 	registerListeners() {
+		this._notebookService.listNotebookEditors().forEach((e) => {
+			this._addNotebookEditor(e);
+		});
+
 		this._register(this._notebookService.onDidChangeActiveEditor(e => {
 			this._proxy.$acceptDocumentAndEditorsDelta({
 				newActiveEditor: e
@@ -185,6 +189,10 @@ export class MainThreadNotebooks extends Disposable implements MainThreadNoteboo
 		this._register(this.accessibilityService.onDidChangeScreenReaderOptimized(() => {
 			updateOrder();
 		}));
+
+		const activeEditorPane = this.editorService.activeEditorPane as any | undefined;
+		const notebookEditor = activeEditorPane?.isNotebookEditor ? activeEditorPane.getControl() : undefined;
+		this._updateState(notebookEditor);
 	}
 
 	async addNotebookDocument(data: INotebookModelAddedData) {
@@ -196,10 +204,14 @@ export class MainThreadNotebooks extends Disposable implements MainThreadNoteboo
 	private _addNotebookEditor(e: IEditor) {
 		this._toDisposeOnEditorRemove.set(e.getId(), combinedDisposable(
 			e.onDidChangeModel(() => this._updateState()),
-			e.onDidFocusEditorWidget(() => this._updateState(e)),
+			e.onDidFocusEditorWidget(() => {
+				this._updateState(e);
+			}),
 		));
 
-		this._updateState();
+		const activeEditorPane = this.editorService.activeEditorPane as any | undefined;
+		const notebookEditor = activeEditorPane?.isNotebookEditor ? activeEditorPane.getControl() : undefined;
+		this._updateState(notebookEditor);
 	}
 
 	private _removeNotebookEditor(e: IEditor) {
@@ -229,6 +241,10 @@ export class MainThreadNotebooks extends Disposable implements MainThreadNoteboo
 			}
 		}
 
+		if (!activeEditor && focusedNotebookEditor) {
+			activeEditor = focusedNotebookEditor.getId();
+		}
+
 		// editors always have view model attached, which means there is already a document in exthost.
 		const newState = new DocumentAndEditorState(documents, editors, activeEditor);
 		const delta = DocumentAndEditorState.compute(this._currentState, newState);
diff --git a/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts b/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts
index bf224dc5e88..7fe5f7fc900 100644
--- a/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts
+++ b/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts
@@ -202,7 +202,6 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditor
 		// Note - focus going to the webview will fire 'blur', but the webview element will be
 		// a descendent of the notebook editor root.
 		this.editorFocus?.set(DOM.isAncestor(document.activeElement, this.overlayContainer));
-		this._onDidFocusEditorWidget.fire();
 	}
 
 	hasFocus() {
@@ -365,6 +364,7 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditor
 	focus() {
 		this.editorFocus?.set(true);
 		this.list?.domFocus();
+		this._onDidFocusEditorWidget.fire();
 	}
 
 	async setModel(textModel: NotebookTextModel, viewState: INotebookEditorViewState | undefined, options: EditorOptions | undefined): Promise<void> {
diff --git a/src/vs/workbench/contrib/notebook/common/notebookEditorModel.ts b/src/vs/workbench/contrib/notebook/common/notebookEditorModel.ts
index d7fa13eae3b..a6784c58295 100644
--- a/src/vs/workbench/contrib/notebook/common/notebookEditorModel.ts
+++ b/src/vs/workbench/contrib/notebook/common/notebookEditorModel.ts
@@ -232,7 +232,7 @@ export class NotebookEditorModelManager extends Disposable implements INotebookE
 		else {
 			// didCreateModel = true;
 			const newModel = model = this.instantiationService.createInstance(NotebookEditorModel, resource, viewType);
-			modelPromise = model.load();
+			modelPromise = model.load({ editorId });
 
 			this.registerModel(newModel);
 		}