diff --git a/src/extension/inlineEdits/vscode-node/features/diagnosticsBasedCompletions/diagnosticsCompletions.ts b/src/extension/inlineEdits/vscode-node/features/diagnosticsBasedCompletions/diagnosticsCompletions.ts
index a3c2c1e20..97562d7b2 100644
--- a/src/extension/inlineEdits/vscode-node/features/diagnosticsBasedCompletions/diagnosticsCompletions.ts
+++ b/src/extension/inlineEdits/vscode-node/features/diagnosticsBasedCompletions/diagnosticsCompletions.ts
@@ -21,6 +21,7 @@ import { Range } from '../../../../../util/vs/editor/common/core/range';
 import { OffsetRange } from '../../../../../util/vs/editor/common/core/ranges/offsetRange';
 import { INextEditDisplayLocation } from '../../../node/nextEditResult';
 import { IVSCodeObservableDocument, IVSCodeObservableNotebookDocument, IVSCodeObservableTextDocument } from '../../parts/vscodeWorkspace';
+import { coalesce } from '../../../../../util/vs/base/common/arrays';
 
 export interface IDiagnosticCodeAction {
 	edit: TextReplacement;
@@ -312,10 +313,21 @@ export class CodeAction {
 	}
 
 	getEditForWorkspaceDocument(workspaceDocument: IVSCodeObservableDocument): TextReplacement[] | undefined {
-		if (!this.edit) {
+		const edit = this.edit;
+		if (!edit) {
 			return undefined;
 		}
-		return this.edit.get(workspaceDocument.id.toUri()).map(toInternalTextEdit);
+		if (workspaceDocument.kind === 'textDocument') {
+			return edit.get(workspaceDocument.id.toUri()).map(e => toInternalTextEdit(e.range, e.newText));
+		} else if (workspaceDocument.kind === 'notebookDocument') {
+			const edits = coalesce(workspaceDocument.notebook.getCells().flatMap(cell => {
+				return edit.get(cell.document.uri).map(e => {
+					const range = workspaceDocument.toRange(cell.document, e.range);
+					return range ? toInternalTextEdit(range, e.newText) : undefined;
+				});
+			}));
+			return edits.length ? edits : undefined;
+		}
 	}
 
 	getDiagnosticsReferencedInCommand(): Diagnostic[] {
@@ -377,8 +389,8 @@ export function toExternalPosition(position: Position): vscode.Position {
 	return new vscode.Position(position.lineNumber - 1, position.column - 1);
 }
 
-export function toInternalTextEdit(edit: vscode.TextEdit): TextReplacement {
-	return new TextReplacement(toInternalRange(edit.range), edit.newText);
+export function toInternalTextEdit(range: vscode.Range, newText: string): TextReplacement {
+	return new TextReplacement(toInternalRange(range), newText);
 }
 
 export function toExternalTextEdit(edit: TextReplacement): vscode.TextEdit {
diff --git a/src/extension/inlineEdits/vscode-node/features/diagnosticsCompletionProcessor.ts b/src/extension/inlineEdits/vscode-node/features/diagnosticsCompletionProcessor.ts
index 3513bccbc..e3bd5fe92 100644
--- a/src/extension/inlineEdits/vscode-node/features/diagnosticsCompletionProcessor.ts
+++ b/src/extension/inlineEdits/vscode-node/features/diagnosticsCompletionProcessor.ts
@@ -38,6 +38,7 @@ import { AnyDiagnosticCompletionItem, AnyDiagnosticCompletionProvider } from './
 import { AsyncDiagnosticCompletionProvider } from './diagnosticsBasedCompletions/asyncDiagnosticsCompletionProvider';
 import { Diagnostic, DiagnosticCompletionItem, DiagnosticInlineEditRequestLogContext, DiagnosticSeverity, distanceToClosestDiagnostic, IDiagnosticCompletionProvider, log, logList, sortDiagnosticsByDistance, toInternalPosition } from './diagnosticsBasedCompletions/diagnosticsCompletions';
 import { ImportDiagnosticCompletionItem, ImportDiagnosticCompletionProvider } from './diagnosticsBasedCompletions/importDiagnosticsCompletionProvider';
+import { isNotebookCell } from '../../../../util/common/notebooks';
 
 interface IDiagnosticsCompletionState<T extends DiagnosticCompletionItem = DiagnosticCompletionItem> {
 	completionItem: T | null;
@@ -226,9 +227,13 @@ export class DiagnosticsCompletionProcessor extends Disposable {
 
 		this._rejectionCollector = new RejectionCollector(this._workspace, s => this._tracer.trace(s));
 
+		const isValidEditor = (editor: vscode.TextEditor | undefined): editor is vscode.TextEditor => {
+			return !!editor && (isNotebookCell(editor.document.uri) || isEditorFromEditorGrid(editor));
+		};
+
 		this._register(this._languageDiagnosticsService.onDidChangeDiagnostics(async e => {
 			const activeEditor = this._tabsAndEditorsService.activeTextEditor;
-			if (!activeEditor || !isEditorFromEditorGrid(activeEditor)) {
+			if (!isValidEditor(activeEditor)) {
 				return;
 			}
 
@@ -241,8 +246,7 @@ export class DiagnosticsCompletionProcessor extends Disposable {
 		}));
 
 		this._register(this._tabsAndEditorsService.onDidChangeActiveTextEditor(async e => {
-			const activeEditor = e;
-			if (!activeEditor || !isEditorFromEditorGrid(activeEditor)) {
+			if (!isValidEditor(e)) {
 				return;
 			}
 
@@ -251,7 +255,7 @@ export class DiagnosticsCompletionProcessor extends Disposable {
 
 		this._register(vscode.window.onDidChangeTextEditorSelection(async e => {
 			const activeEditor = this._tabsAndEditorsService.activeTextEditor;
-			if (!activeEditor || !isEditorFromEditorGrid(activeEditor)) {
+			if (!isValidEditor(activeEditor)) {
 				return;
 			}
 
@@ -291,9 +295,14 @@ export class DiagnosticsCompletionProcessor extends Disposable {
 		const workspaceDocument = this._workspace.getDocumentByTextDocument(activeTextEditor.document);
 		if (!workspaceDocument) { return; }
 
-		const log = new DiagnosticInlineEditRequestLogContext();
+		const range = new vscode.Range(activeTextEditor.selection.active, activeTextEditor.selection.active);
+		const selection = workspaceDocument.toRange(activeTextEditor.document, range);
+		if (!selection) {
+			return;
+		}
 
-		const cursor = toInternalPosition(activeTextEditor.selection.active);
+		const cursor = toInternalPosition(selection.start);
+		const log = new DiagnosticInlineEditRequestLogContext();
 
 		const { availableDiagnostics, relevantDiagnostics } = this._getDiagnostics(workspaceDocument, cursor, log);
 		const diagnosticsSorted = sortDiagnosticsByDistance(relevantDiagnostics, cursor);
diff --git a/src/extension/inlineEdits/vscode-node/parts/vscodeWorkspace.ts b/src/extension/inlineEdits/vscode-node/parts/vscodeWorkspace.ts
index 136315841..6a7c6908f 100644
--- a/src/extension/inlineEdits/vscode-node/parts/vscodeWorkspace.ts
+++ b/src/extension/inlineEdits/vscode-node/parts/vscodeWorkspace.ts
@@ -387,6 +387,7 @@ export interface IVSCodeObservableTextDocument extends IObservableDocument {
 	fromOffsetRange(textDocument: TextDocument, range: OffsetRange): Range | undefined;
 	fromRange(textDocument: TextDocument, range: Range): Range | undefined;
 	toOffsetRange(textDocument: TextDocument, range: Range): OffsetRange | undefined;
+	toRange(textDocument: TextDocument, range: Range): Range | undefined;
 }
 
 abstract class AbstractVSCodeObservableDocument {
@@ -445,6 +446,10 @@ class VSCodeObservableTextDocument extends AbstractVSCodeObservableDocument impl
 		return new OffsetRange(textDocument.offsetAt(range.start), textDocument.offsetAt(range.end));
 	}
 
+	toRange(_textDocument: TextDocument, range: Range): Range | undefined {
+		return range;
+	}
+
 	fromRange(_textDocument: TextDocument, range: Range): Range | undefined {
 		return range;
 	}
@@ -466,6 +471,7 @@ export interface IVSCodeObservableNotebookDocument extends IObservableDocument {
 	fromRange(textDocument: TextDocument, range: Range): Range | undefined;
 	fromRange(range: Range): [TextDocument, Range][];
 	projectDiagnostics(cell: TextDocument, diagnostics: readonly Diagnostic[]): Diagnostic[];
+	toRange(textDocument: TextDocument, range: Range): Range | undefined;
 }
 
 class VSCodeObservableNotebookDocument extends AbstractVSCodeObservableDocument implements IVSCodeObservableNotebookDocument {
@@ -523,6 +529,14 @@ class VSCodeObservableNotebookDocument extends AbstractVSCodeObservableDocument
 		}
 		return toAltDiagnostics(this.altNotebook, cell, diagnostics);
 	}
+	toRange(textDocument: TextDocument, range: Range): Range | undefined {
+		const cell = this.altNotebook.getCell(textDocument);
+		if (!cell) {
+			return undefined;
+		}
+		const ranges = this.altNotebook.toAltRange(cell, [range]);
+		return ranges.length > 0 ? ranges[0] : undefined;
+	}
 }
 
 export type IVSCodeObservableDocument = IVSCodeObservableTextDocument | IVSCodeObservableNotebookDocument;
@@ -691,4 +705,3 @@ export function editFromTextDocumentContentChangeEvents(events: readonly TextDoc
 	const replacementsInApplicationOrder = events.map(e => StringReplacement.replace(OffsetRange.ofStartAndLength(e.rangeOffset, e.rangeLength), e.text));
 	return StringEdit.composeSequentialReplacements(replacementsInApplicationOrder);
 }
-