[AI Assistance] Remove inspector protocol changes after save

We want to minimize the layout/style shift, so we apply things
to the working copy and then remove the inspector stylesheets

Fixed: 407972738
Change-Id: Ie6700ac485ccc8c409506a41051d05f29b8f7977
Reviewed-on: https://chromium-review.googlesource.com/c/devtools/devtools-frontend/+/6434111
Commit-Queue: Alex Rudenko <[email protected]>
Reviewed-by: Alex Rudenko <[email protected]>
Auto-Submit: Nikolay Vitkov <[email protected]>

Author: Lightning00Blade

front_end/panels/ai_assistance/PatchWidget.test.ts

@@ -281,8 +281,6 @@ Files:
     let uiSourceCode: Workspace.UISourceCode.UISourceCode;
     let commitWorkingCopyStub:
         sinon.SinonStub<Parameters<typeof Workspace.UISourceCode.UISourceCode.prototype.commitWorkingCopy>>;
-    let resetWorkingCopyStub:
-        sinon.SinonStub<Parameters<typeof Workspace.UISourceCode.UISourceCode.prototype.resetWorkingCopy>>;
 
     beforeEach(() => {
       uiSourceCode = createTestFilesystem('file://test').uiSourceCode;
@@ -296,49 +294,50 @@ Files:
 
       commitWorkingCopyStub =
           sinon.stub(Workspace.UISourceCode.UISourceCode.prototype, 'commitWorkingCopy').callThrough();
-      resetWorkingCopyStub =
-          sinon.stub(Workspace.UISourceCode.UISourceCode.prototype, 'resetWorkingCopy').callThrough();
     });
 
-    it('on apply should call handle function and stash changes', async () => {
+    it('on apply should should not call stashChanges', async () => {
       const {
         view,
         widget,
-      } = await createPatchWidget({aidaClient: mockAidaClient([[{explanation: 'patch applied'}]])});
+      } = await createPatchWidget({
+        aidaClient: mockAidaClient([[{explanation: 'patch applied'}]]),
+      });
       widget.changeSummary = 'body { background-color: red; }';
       const changeManager = sinon.createStubInstance(AiAssistanceModel.ChangeManager);
+      changeManager.stashChanges.returns(Promise.resolve());
       widget.changeManager = changeManager;
       view.input.onApplyToWorkspace();
       await view.nextInput;
-      assert.isTrue(changeManager.stashChanges.calledOnce);
+      assert.isTrue(changeManager.stashChanges.notCalled);
     });
 
-    it('save all should commit the working copy of the changed UI codes to the disk and render savedToDisk view',
-       async () => {
-         const {view, widget} = await createPatchWidgetWithDiffView();
-         const changeManager = sinon.createStubInstance(AiAssistanceModel.ChangeManager);
-         widget.changeManager = changeManager;
-         uiSourceCode.setWorkingCopy('working copy');
+    it('on save should stash changes', async () => {
+      const {view, widget} = await createPatchWidgetWithDiffView();
+      const changeManager = sinon.createStubInstance(AiAssistanceModel.ChangeManager);
+      changeManager.stashChanges.returns(Promise.resolve());
+      widget.changeManager = changeManager;
+      uiSourceCode.setWorkingCopy('working copy');
 
-         view.input.onSaveAll();
-         const nextInput = await view.nextInput;
+      view.input.onSaveAll();
+      const nextInput = await view.nextInput;
 
-         assert.isTrue(nextInput.savedToDisk);
-         assert.isTrue(commitWorkingCopyStub.called, 'Expected commitWorkingCopy to be called but it is not called');
-         assert.isTrue(changeManager.dropStashedChanges.calledOnce);
-       });
+      assert.isTrue(nextInput.savedToDisk);
+      assert.isTrue(commitWorkingCopyStub.called, 'Expected commitWorkingCopy to be called but it is not called');
+      assert.isTrue(changeManager.stashChanges.calledOnce);
+    });
 
     it('discard should discard the working copy and render the view without patchSuggestion', async () => {
       const {view, widget} = await createPatchWidgetWithDiffView();
       const changeManager = sinon.createStubInstance(AiAssistanceModel.ChangeManager);
+      changeManager.stashChanges.returns(Promise.resolve());
       widget.changeManager = changeManager;
       uiSourceCode.setWorkingCopy('working copy');
 
       view.input.onDiscard();
       const nextInput = await view.nextInput;
 
       assert.strictEqual(nextInput.patchSuggestionState, AiAssistance.PatchWidget.PatchSuggestionState.INITIAL);
-      assert.isTrue(resetWorkingCopyStub.called, 'Expected resetWorkingCopy to be called but it is not called');
       assert.isTrue(changeManager.popStashedChanges.calledOnce);
     });
   });

front_end/panels/ai_assistance/PatchWidget.ts

@@ -541,7 +541,6 @@ export class PatchWidget extends UI.Widget.Widget {
     this.requestUpdate();
     const {response, processedFiles} = await this.#applyPatch(changeSummary);
     if (response?.type === AiAssistanceModel.ResponseType.ANSWER) {
-      await this.changeManager?.stashChanges();
       this.#patchSuggestionState = PatchSuggestionState.SUCCESS;
     } else if (
         response?.type === AiAssistanceModel.ResponseType.ERROR &&
@@ -570,11 +569,15 @@ ${processedFiles.map(filename => `* ${filename}`).join('\n')}`;
 
   #onSaveAll(): void {
     this.#workspaceDiff.modifiedUISourceCodes().forEach(modifiedUISourceCode => {
-      modifiedUISourceCode.commitWorkingCopy();
+      if (!modifiedUISourceCode.url().startsWith('inspector://')) {
+        modifiedUISourceCode.commitWorkingCopy();
+      }
+    });
+    void this.changeManager?.stashChanges().then(() => {
+      this.changeManager?.dropStashedChanges();
     });
 
     this.#savedToDisk = true;
-    this.changeManager?.dropStashedChanges();
     this.requestUpdate();
   }