diff --git a/src/sql/parts/notebook/models/notebookModel.ts b/src/sql/parts/notebook/models/notebookModel.ts
index b575303c57..8a3e521070 100644
--- a/src/sql/parts/notebook/models/notebookModel.ts
+++ b/src/sql/parts/notebook/models/notebookModel.ts
@@ -11,7 +11,6 @@ import { localize } from 'vs/nls';
 import { Event, Emitter } from 'vs/base/common/event';
 import { Disposable, IDisposable } from 'vs/base/common/lifecycle';
 
-import { CellModel } from './cell';
 import { IClientSession, INotebookModel, IDefaultConnection, INotebookModelOptions, ICellModel, NotebookContentChange, notebookConstants } from './modelInterfaces';
 import { NotebookChangeType, CellType } from 'sql/parts/notebook/models/contracts';
 import { nbversion } from '../notebookConstants';
diff --git a/src/sql/parts/notebook/notebook.component.ts b/src/sql/parts/notebook/notebook.component.ts
index c25c38b024..3d77252fd5 100644
--- a/src/sql/parts/notebook/notebook.component.ts
+++ b/src/sql/parts/notebook/notebook.component.ts
@@ -176,7 +176,8 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 		this._model.cells.forEach(cell => {
 			cell.trustedMode = isTrusted;
 		});
-		//TODO: Handle dirty for trust?
+		//Updates dirty state
+		this._notebookParams.input && this._notebookParams.input.updateModel();
 		this.detectChanges();
 	}
 
diff --git a/src/sql/parts/notebook/notebookInput.ts b/src/sql/parts/notebook/notebookInput.ts
index dea149e9ba..3ba8a2faef 100644
--- a/src/sql/parts/notebook/notebookInput.ts
+++ b/src/sql/parts/notebook/notebookInput.ts
@@ -138,7 +138,6 @@ export class NotebookInput extends EditorInput {
 	private _parentContainer: HTMLElement;
 	private readonly _layoutChanged: Emitter<void> = this._register(new Emitter<void>());
 	private _model: NotebookEditorModel;
-	private _untitledEditorService: IUntitledEditorService;
 	private _contentManager: IContentManager;
 	private _providersLoaded: Promise<void>;
 
@@ -146,13 +145,11 @@ export class NotebookInput extends EditorInput {
 		private resource: URI,
 		private _textInput: UntitledEditorInput,
 		@ITextModelService private textModelService: ITextModelService,
-		@IUntitledEditorService untitledEditorService: IUntitledEditorService,
 		@IInstantiationService private instantiationService: IInstantiationService,
 		@INotebookService private notebookService: INotebookService,
 		@IExtensionService private extensionService: IExtensionService
 	) {
 		super();
-		this._untitledEditorService = untitledEditorService;
 		this.resource = resource;
 		this._standardKernels = [];
 		this._providersLoaded = this.assignProviders();