fix: edge cases

Author: yaslam-dev

ext-src/commands/DiscardSuggestion.ts

@@ -1,6 +1,6 @@
 import * as vscode from 'vscode';
 import { CurrentQuestion, Analyze, Session, AnalyseMetaData } from '../state';
-import { GenerateDecorations } from '../helpers/GenerateDecorations';
+import { GenerateDecorations, decorationType } from '../helpers/GenerateDecorations';
 import { FeedbackSuggestionPayload, feedbackService } from '../services';
 import CONSTANTS from '../constants';
 import Utils from '../utils';
@@ -45,6 +45,8 @@ export function activateDiscardCommand(context: vscode.ExtensionContext): void {
       return;
     }
 
+    const copyProblems = { ...problems };
+
     const payload: FeedbackSuggestionPayload = {
       problemId,
       discarded: true,
@@ -57,40 +59,30 @@ export function activateDiscardCommand(context: vscode.ExtensionContext): void {
       _debug?.appendLine(err.message);
     }
 
-    problems[key] = {
-      ...problems[key],
-      isDiscarded: true,
-    };
+    copyProblems[key].isDiscarded = true;
 
     try {
-      await analyzeState.set(problems);
+      analyzeState.set(copyProblems).then(() => {
+        const results: AnalyseMetaData[] = [];
+
+        for (const [, value] of Object.entries(copyProblems)) {
+          if (!value.isDiscarded) {
+            results.push(value);
+          }
+        }
+
+        const { decorations } = GenerateDecorations(results, editor);
+        editor.setDecorations(decorationType, []);
+        editor.setDecorations(decorationType, decorations);
+        currentQuestion.clear();
+
+        return;
+      });
     } catch (error: any) {
       _debug.appendLine(error);
 
       return;
     }
-
-    const analyzeValue = analyzeState.get()?.value;
-
-    if (!analyzeValue) {
-      _debug?.appendLine('Metabob: analyzeValue is undefined in Discard Suggestion');
-
-      return;
-    }
-
-    // Map all non discarded problems to the results array
-    const results: AnalyseMetaData[] = Object.keys(analyzeValue)
-      .filter(localKey => !analyzeValue[key].isDiscarded || localKey === key)
-      .map(key => analyzeValue[key]);
-
-    const decorations = GenerateDecorations(results, editor);
-    editor.setDecorations(decorations.decorationType, []);
-    editor.setDecorations(decorations.decorationType, decorations.decorations);
-    currentQuestion.clear();
-
-    // TODO: Send current question update to webview
-
-    return;
   };
 
   context.subscriptions.push(vscode.commands.registerCommand(command, commandHandler));

ext-src/events.ts

@@ -1,7 +1,7 @@
 import { EventEmitter } from 'vscode';
 
 export interface AnalysisEvents {
-  type: 'Analysis_Error' | 'Analysis_Completed' | 'FIX_SUGGESTION';
+  type: 'Analysis_Error' | 'Analysis_Completed' | 'Analysis_Called_On_Save' | 'FIX_SUGGESTION';
   data: any;
 }
 

ext-src/extension.ts

@@ -74,23 +74,6 @@ export function activate(context: vscode.ExtensionContext): void {
     debugChannel.appendLine(`Metabob: ${error}`);
   }
 
-  // Analyze on Save functionality is only ran if the user enabled it.
-  if (analyzeDocumentOnSaveConfig && analyzeDocumentOnSaveConfig === true) {
-    context.subscriptions.push(
-      vscode.workspace.onDidSaveTextDocument(document => {
-        // Will check if the current document is valid code file.
-        if (Util.isValidDocument(document)) {
-          AnalyzeDocumentOnSave(
-            {
-              document,
-            },
-            context,
-          );
-        }
-      }),
-    );
-  }
-
   // If the user changes the global config from CMD + Shift + P -> User Setting -> Metabob
   // Then, we would update global state or reload the application
   context.subscriptions.push(
@@ -143,6 +126,91 @@ export function activate(context: vscode.ExtensionContext): void {
 
   const extensionEventEmitter = getExtensionEventEmitter();
 
+  // Analyze on Save functionality is only ran if the user enabled it.
+  context.subscriptions.push(
+    vscode.workspace.onDidSaveTextDocument(document => {
+      // Will check if the current document is valid code file.
+      if (analyzeDocumentOnSaveConfig && analyzeDocumentOnSaveConfig === true) {
+        if (!Util.isValidDocument) {
+          return;
+        }
+
+        AnalyzeDocumentOnSave(
+          {
+            document,
+          },
+          context,
+        );
+        extensionEventEmitter.fire({
+          type: 'Analysis_Called_On_Save',
+          data: {},
+        });
+      }
+    }),
+  );
+
+  context.subscriptions.push(
+    vscode.workspace.onDidCloseTextDocument(() => {
+      extensionEventEmitter.fire({
+        type: 'Analysis_Completed',
+        data: {},
+      });
+    }),
+  );
+  context.subscriptions.push(
+    vscode.workspace.onDidOpenTextDocument(() => {
+      const currentEditor = vscode.window.activeTextEditor;
+      if (!currentEditor) return;
+
+      if (!Util.isValidDocument(currentEditor.document)) {
+        return;
+      }
+
+      const documentMetaData = Util.extractMetaDataFromDocument(currentEditor.document);
+
+      const filename: string | undefined = documentMetaData.relativePath.split('/').pop();
+
+      if (!filename) return;
+
+      const analyzeState = new Analyze(context);
+
+      const analyzeValue = analyzeState.get()?.value;
+
+      if (!analyzeValue) return;
+
+      const results: AnalyseMetaData[] = [];
+      const analzeResultsOfCurrentFile: AnalyzeState = {};
+
+      for (const [key, value] of Object.entries(analyzeValue)) {
+        const splitString: string | undefined = key.split('@@')[0];
+        if (splitString === undefined) continue;
+
+        if (splitString === filename && value.isDiscarded === false) {
+          results.push(value);
+
+          analzeResultsOfCurrentFile[key] = value;
+        }
+      }
+
+      if (results.length === 0) {
+        extensionEventEmitter.fire({
+          type: 'Analysis_Completed',
+          data: analzeResultsOfCurrentFile,
+        });
+
+        return;
+      }
+
+      const { decorations } = GenerateDecorations(results, currentEditor);
+
+      currentEditor.setDecorations(decorationType, []);
+      currentEditor.setDecorations(decorationType, decorations);
+      extensionEventEmitter.fire({
+        type: 'Analysis_Completed',
+        data: { ...analzeResultsOfCurrentFile },
+      });
+    }),
+  );
   context.subscriptions.push(
     vscode.window.onDidChangeActiveTextEditor(currentEditor => {
       if (!currentEditor) return;

ext-src/providers/recommendation.provider.ts

@@ -16,9 +16,9 @@ import {
 } from 'vscode';
 import { Configuration, CreateChatCompletionRequest, OpenAIApi } from 'openai';
 import { explainService, ExplainProblemPayload, SuggestRecomendationPayload } from '../services';
-import { CurrentQuestion, CurrentQuestionState, Session } from '../state';
+import { CurrentQuestion, CurrentQuestionState, Session, Analyze, AnalyseMetaData } from '../state';
 import { BackendService, GetChatGPTToken } from '../config';
-import { GenerateDecorations } from '../helpers';
+import { GenerateDecorations, decorationType } from '../helpers';
 import { DiscardCommandHandler, EndorseCommandHandler } from '../commands';
 import CONSTANTS from '../constants';
 import Util from '../utils';
@@ -31,6 +31,7 @@ export class RecommendationWebView implements WebviewViewProvider {
   private readonly extensionURI: Uri;
   private readonly extensionContext: ExtensionContext;
   private extensionEventEmitter: EventEmitter<AnalysisEvents>;
+  private eventEmitterQueue: Array<AnalysisEvents> = [];
 
   constructor(
     extensionPath: string,
@@ -86,6 +87,29 @@ export class RecommendationWebView implements WebviewViewProvider {
         return;
       }
 
+      let interval: NodeJS.Timeout | undefined = undefined;
+
+      if (!this._view.visible) {
+        this.eventEmitterQueue.push(event);
+        interval = setInterval(() => {
+          if (this?._view === null || this?._view === undefined || !this?._view.webview) {
+            return;
+          }
+
+          if (!this._view.visible) {
+            return;
+          }
+
+          const latestEvent = this.eventEmitterQueue.pop();
+          if (latestEvent) {
+            this?._view?.webview?.postMessage(latestEvent);
+          }
+
+          this.eventEmitterQueue = [];
+          clearInterval(interval);
+        }, 500);
+      }
+
       this._view.webview.postMessage(event);
     });
   }
@@ -286,7 +310,7 @@ export class RecommendationWebView implements WebviewViewProvider {
     const comment = `\t\t#${input}`;
     const position = new Position(startLine - 1, 0); // convert line number to position
     editor.edit((editBuilder: TextEditorEdit) => {
-      editBuilder.insert(position, comment + '\n');
+      editBuilder.replace(position, comment + '\n');
     });
   }
 
@@ -326,18 +350,56 @@ export class RecommendationWebView implements WebviewViewProvider {
       throw new Error('handleApplyRecommendation: Editor or Init Data is undefined');
     }
 
-    const startLine = initData.vuln.startLine;
-    const endLine = initData.vuln.endLine;
-    const comment = `${input.replace('```', '')}`;
+    debugChannel.appendLine('here1');
+    const setAnalyzeState = new Analyze(this.extensionContext);
+    const getanalyzeState = new Analyze(this.extensionContext).get()?.value;
 
-    const data = initData.vuln;
-    const start = new Position(startLine - 1, 0); // convert line number to position
-    const end = new Position(endLine, 0); // convert line number to position
-    const range = new Range(start, end);
-    editor.edit((editBuilder: TextEditorEdit) => {
-      const decorations = GenerateDecorations([{ ...data }], editor);
-      editor.setDecorations(decorations.decorationType, []);
-      editBuilder.replace(range, comment + '\n');
+    debugChannel.appendLine('here2');
+
+    if (!getanalyzeState) {
+      debugChannel.appendLine('here3');
+      throw new Error('Analze is undefined');
+    }
+    debugChannel.appendLine('here4');
+
+    const copyAnalyzeValue = { ...getanalyzeState };
+
+    const key = `${initData.path}@@${initData.id}`;
+    copyAnalyzeValue[key].isDiscarded = true;
+
+    const results: AnalyseMetaData[] = [];
+
+    debugChannel.appendLine('here5');
+
+    for (const [, value] of Object.entries(copyAnalyzeValue)) {
+      if (!value.isDiscarded) {
+        results.push(value);
+      }
+    }
+
+    debugChannel.appendLine('here6');
+
+    setAnalyzeState.set({ ...copyAnalyzeValue }).then(() => {
+      debugChannel.appendLine('here7');
+
+      const startLine = initData.vuln.startLine;
+      const endLine = initData.vuln.endLine;
+      const comment = `${input.replace('```', '')}`;
+
+      const start = new Position(startLine - 1, 0); // convert line number to position
+      const end = new Position(endLine, 0); // convert line number to position
+      const range = new Range(start, end);
+
+      const { decorations } = GenerateDecorations(results, editor);
+
+      editor.setDecorations(decorationType, []);
+      editor.setDecorations(decorationType, decorations);
+
+      debugChannel.appendLine('here8');
+
+      editor.edit((editBuilder: TextEditorEdit) => {
+        editBuilder.replace(range, comment + '\n');
+      });
     });
   }
 
@@ -466,7 +528,9 @@ export class RecommendationWebView implements WebviewViewProvider {
           }
           try {
             this.handleApplyRecommendation(input, initData);
-          } catch {
+          } catch (error: any) {
+            debugChannel.appendLine(`Metabob: Apply Recommendation Error ${JSON.stringify(error)}`);
+
             this._view.webview.postMessage({
               type: 'applyRecommendation:Error',
               data: {},

ext-src/utils.ts

@@ -118,4 +118,12 @@ export default class Utils {
 
     return decor;
   }
+
+  static openFileInNewTab(filePath: string): Thenable<vscode.TextEditor | undefined> {
+    const uri = vscode.Uri.file(filePath);
+
+    return vscode.workspace.openTextDocument(uri).then(document => {
+      return vscode.window.showTextDocument(document, vscode.ViewColumn.One);
+    });
+  }
 }