diff --git a/src/vs/workbench/contrib/notebook/browser/view/cellParts/codeCell.ts b/src/vs/workbench/contrib/notebook/browser/view/cellParts/codeCell.ts
index 78fd5c0bc4a..2ec0ceb0bef 100644
--- a/src/vs/workbench/contrib/notebook/browser/view/cellParts/codeCell.ts
+++ b/src/vs/workbench/contrib/notebook/browser/view/cellParts/codeCell.ts
@@ -69,10 +69,8 @@ export class CodeCell extends Disposable {
 		this.registerDecorations();
 		this.registerMouseListener();
 
-		this._register(notebookExecutionStateService.onDidChangeCellExecution(e => {
-			if (e.affectsCell(this.viewCell.uri)) {
-				this.cellParts.updateForExecutionState(this.viewCell, e);
-			}
+		this._register(Event.any(this.viewCell.onDidStartExecution, this.viewCell.onDidStopExecution)((e) => {
+			this.cellParts.updateForExecutionState(this.viewCell, e);
 		}));
 
 		this._register(this.viewCell.onDidChangeState(e => {
diff --git a/src/vs/workbench/contrib/notebook/browser/viewModel/codeCellViewModel.ts b/src/vs/workbench/contrib/notebook/browser/viewModel/codeCellViewModel.ts
index a37a878645a..3858986efe8 100644
--- a/src/vs/workbench/contrib/notebook/browser/viewModel/codeCellViewModel.ts
+++ b/src/vs/workbench/contrib/notebook/browser/viewModel/codeCellViewModel.ts
@@ -21,7 +21,7 @@ import { NotebookOptionsChangeEvent } from 'vs/workbench/contrib/notebook/browse
 import { INotebookService } from 'vs/workbench/contrib/notebook/common/notebookService';
 import { BaseCellViewModel } from './baseCellViewModel';
 import { NotebookLayoutInfo } from 'vs/workbench/contrib/notebook/browser/notebookViewEvents';
-import { INotebookExecutionStateService } from 'vs/workbench/contrib/notebook/common/notebookExecutionStateService';
+import { ICellExecutionStateChangedEvent } from 'vs/workbench/contrib/notebook/common/notebookExecutionStateService';
 
 export class CodeCellViewModel extends BaseCellViewModel implements ICellViewModel {
 	readonly cellKind = CellKind.Code;
@@ -29,9 +29,9 @@ export class CodeCellViewModel extends BaseCellViewModel implements ICellViewMod
 	protected readonly _onLayoutInfoRead = this._register(new Emitter<void>());
 	readonly onLayoutInfoRead = this._onLayoutInfoRead.event;
 
-	protected readonly _onDidStartExecution = this._register(new Emitter<void>());
+	protected readonly _onDidStartExecution = this._register(new Emitter<ICellExecutionStateChangedEvent>());
 	readonly onDidStartExecution = this._onDidStartExecution.event;
-	protected readonly _onDidStopExecution = this._register(new Emitter<void>());
+	protected readonly _onDidStopExecution = this._register(new Emitter<ICellExecutionStateChangedEvent>());
 	readonly onDidStopExecution = this._onDidStopExecution.event;
 
 	protected readonly _onDidChangeOutputs = this._register(new Emitter<NotebookCellOutputsSplice>());
@@ -125,7 +125,6 @@ export class CodeCellViewModel extends BaseCellViewModel implements ICellViewMod
 		readonly viewContext: ViewContext,
 		@IConfigurationService configurationService: IConfigurationService,
 		@INotebookService private readonly _notebookService: INotebookService,
-		@INotebookExecutionStateService private readonly _notebookExecutionStateService: INotebookExecutionStateService,
 		@ITextModelService modelService: ITextModelService,
 		@IUndoRedoService undoRedoService: IUndoRedoService,
 		@ICodeEditorService codeEditorService: ICodeEditorService
@@ -154,16 +153,6 @@ export class CodeCellViewModel extends BaseCellViewModel implements ICellViewMod
 			dispose(removedOutputs);
 		}));
 
-		this._register(this._notebookExecutionStateService.onDidChangeCellExecution(e => {
-			if (e.affectsCell(model.uri)) {
-				if (e.changed) {
-					this._onDidStartExecution.fire();
-				} else {
-					this._onDidStopExecution.fire();
-				}
-			}
-		}));
-
 		this._outputCollection = new Array(this.model.outputs.length);
 
 		this._layoutInfo = {
@@ -187,6 +176,14 @@ export class CodeCellViewModel extends BaseCellViewModel implements ICellViewMod
 		};
 	}
 
+	updateExecutionState(e: ICellExecutionStateChangedEvent) {
+		if (e.changed) {
+			this._onDidStartExecution.fire(e);
+		} else {
+			this._onDidStopExecution.fire(e);
+		}
+	}
+
 	updateOptions(e: NotebookOptionsChangeEvent) {
 		if (e.cellStatusBarVisibility || e.insertToolbarPosition || e.cellToolbarLocation) {
 			this.layoutChange({});
diff --git a/src/vs/workbench/contrib/notebook/browser/viewModel/notebookViewModelImpl.ts b/src/vs/workbench/contrib/notebook/browser/viewModel/notebookViewModelImpl.ts
index 9308ddf9cf0..d214e5793c7 100644
--- a/src/vs/workbench/contrib/notebook/browser/viewModel/notebookViewModelImpl.ts
+++ b/src/vs/workbench/contrib/notebook/browser/viewModel/notebookViewModelImpl.ts
@@ -33,6 +33,7 @@ import { CellKind, ICell, INotebookSearchOptions, ISelectionState, NotebookCells
 import { cellIndexesToRanges, cellRangesToIndexes, ICellRange, reduceCellRanges } from 'vs/workbench/contrib/notebook/common/notebookRange';
 import { NotebookLayoutInfo, NotebookMetadataChangedEvent } from 'vs/workbench/contrib/notebook/browser/notebookViewEvents';
 import { CellFindMatchModel } from 'vs/workbench/contrib/notebook/browser/contrib/find/findModel';
+import { INotebookExecutionStateService } from 'vs/workbench/contrib/notebook/common/notebookExecutionStateService';
 
 const invalidFunc = () => { throw new Error(`Invalid change accessor`); };
 
@@ -196,6 +197,7 @@ export class NotebookViewModel extends Disposable implements EditorFoldingStateD
 		@IBulkEditService private readonly _bulkEditService: IBulkEditService,
 		@IUndoRedoService private readonly _undoService: IUndoRedoService,
 		@ITextModelService private readonly _textModelService: ITextModelService,
+		@INotebookExecutionStateService notebookExecutionStateService: INotebookExecutionStateService,
 	) {
 		super();
 
@@ -318,6 +320,13 @@ export class NotebookViewModel extends Disposable implements EditorFoldingStateD
 			}
 		}));
 
+		this._register(notebookExecutionStateService.onDidChangeCellExecution(e => {
+			const cell = this.getCellByHandle(e.cellHandle);
+
+			if (cell instanceof CodeCellViewModel) {
+				cell.updateExecutionState(e);
+			}
+		}));
 
 		this._register(this._selectionCollection.onDidChangeSelection(e => {
 			this._onDidChangeSelection.fire(e);
diff --git a/src/vs/workbench/contrib/notebook/test/browser/notebookViewModel.test.ts b/src/vs/workbench/contrib/notebook/test/browser/notebookViewModel.test.ts
index 28c8a31c8e4..b6f13158a18 100644
--- a/src/vs/workbench/contrib/notebook/test/browser/notebookViewModel.test.ts
+++ b/src/vs/workbench/contrib/notebook/test/browser/notebookViewModel.test.ts
@@ -37,6 +37,7 @@ suite('NotebookViewModel', () => {
 	let undoRedoService: IUndoRedoService;
 	let modelService: IModelService;
 	let languageService: ILanguageService;
+	let notebookExecutionStateService: INotebookExecutionStateService;
 
 	suiteSetup(() => {
 		disposables = new DisposableStore();
@@ -46,6 +47,7 @@ suite('NotebookViewModel', () => {
 		undoRedoService = instantiationService.get(IUndoRedoService);
 		modelService = instantiationService.get(IModelService);
 		languageService = instantiationService.get(ILanguageService);
+		notebookExecutionStateService = instantiationService.get(INotebookExecutionStateService);
 
 		instantiationService.stub(IConfigurationService, new TestConfigurationService());
 		instantiationService.stub(IThemeService, new TestThemeService());
@@ -57,7 +59,7 @@ suite('NotebookViewModel', () => {
 		const notebook = new NotebookTextModel('notebook', URI.parse('test'), [], {}, { transientCellMetadata: {}, transientDocumentMetadata: {}, transientOutputs: false, cellContentMetadata: {} }, undoRedoService, modelService, languageService);
 		const model = new NotebookEditorTestModel(notebook);
 		const viewContext = new ViewContext(new NotebookOptions(instantiationService.get(IConfigurationService), instantiationService.get(INotebookExecutionStateService)), new NotebookEventDispatcher(), () => ({} as IBaseCellEditorOptions));
-		const viewModel = new NotebookViewModel('notebook', model.notebook, viewContext, null, { isReadOnly: false }, instantiationService, bulkEditService, undoRedoService, textModelService);
+		const viewModel = new NotebookViewModel('notebook', model.notebook, viewContext, null, { isReadOnly: false }, instantiationService, bulkEditService, undoRedoService, textModelService, notebookExecutionStateService);
 		assert.strictEqual(viewModel.viewType, 'notebook');
 	});