Improve note status marking behavior

src/extension.ts

@@ -181,43 +181,62 @@ export function activate(context: vscode.ExtensionContext) {
     ),
   );
 
-  // set note status as vulnerable button
+  /**
+   * Handles the common logic for setting a note's status via a command.
+   *
+   * @param reply The argument passed by the command (either CommentReply or just the thread).
+   * @param status The NoteStatus to set (TODO, Vulnerable, Not Vulnerable).
+   * @param noteMap The object storing all notes in memory.
+   * @param remoteDb Remote db for collaboration.
+   */
+  const handleSetStatusAction = (
+    reply: vscode.CommentReply | { thread: vscode.CommentThread },
+    status: NoteStatus,
+    noteMap: Map<string, vscode.CommentThread>,
+    remoteDb?: RemoteDb
+  ) => {
+    const thread = reply.thread;
+    // Extract the text of the reply box
+    const text = 'text' in reply ? reply.text : undefined;
+
+    // Set the status (this function handles adding the status change comment)
+    setNoteStatus(
+      thread,
+      status, // New status to set
+      noteMap,
+      '',
+      remoteDb,
+      text // Reply text
+    );
+  };
+
+  // --- Register the status commands ---
+
+  // Set note status as Vulnerable button
   context.subscriptions.push(
     vscode.commands.registerCommand(
       'security-notes.setNoteStatusVulnerable',
-      (commentReply: vscode.CommentReply) =>
-        setNoteStatus(
-          commentReply.thread,
-          NoteStatus.Vulnerable,
-          noteMap,
-          '',
-          remoteDb,
-        ),
-    ),
+      (reply: vscode.CommentReply | { thread: vscode.CommentThread }) =>
+        handleSetStatusAction(reply, NoteStatus.Vulnerable, noteMap, remoteDb)
+    )
   );
 
-  // set note status as not vulnerable button
+  // Set note status as Not Vulnerable button
   context.subscriptions.push(
     vscode.commands.registerCommand(
       'security-notes.setNoteStatusNotVulnerable',
-      (commentReply: vscode.CommentReply) =>
-        setNoteStatus(
-          commentReply.thread,
-          NoteStatus.NotVulnerable,
-          noteMap,
-          '',
-          remoteDb,
-        ),
-    ),
+      (reply: vscode.CommentReply | { thread: vscode.CommentThread }) =>
+        handleSetStatusAction(reply, NoteStatus.NotVulnerable, noteMap, remoteDb)
+    )
   );
 
-  // set note status as TODO button
+  // Set note status as TODO button
   context.subscriptions.push(
     vscode.commands.registerCommand(
       'security-notes.setNoteStatusToDo',
-      (commentReply: vscode.CommentReply) =>
-        setNoteStatus(commentReply.thread, NoteStatus.TODO, noteMap, '', remoteDb),
-    ),
+      (reply: vscode.CommentReply | { thread: vscode.CommentThread }) =>
+        handleSetStatusAction(reply, NoteStatus.TODO, noteMap, remoteDb)
+    )
   );
 
   // webview for importing tool results

src/helpers.ts

@@ -43,26 +43,29 @@ export const setNoteStatus = (
   noteMap: Map<string, vscode.CommentThread>,
   author?: string,
   remoteDb?: RemoteDb,
+  replyText?: string,
 ) => {
   const comment: vscode.Comment | any = thread.comments[0];
 
-  // Remove previous status if any
-  let removed = false;
-  Object.values(NoteStatus).forEach((noteStatus) => {
-    if (!removed && comment.body.toString().startsWith(`[${noteStatus}] `)) {
-      comment.body = comment.body.toString().slice(noteStatus.length + 3);
-      removed = true;
-    }
-  });
+  let originalText = comment.body.toString();
 
-  // Prepend new status on first note comment
-  comment.body = `[${status}] ${comment.body}`;
-  comment.savedBody = comment.body;
+  // Clean up any existing status badges
+  const statusValuesPattern = Object.values(NoteStatus).join('|');
+  const statusRegex = new RegExp(`^\\[(${statusValuesPattern})\\] `, 'g');
+  originalText = originalText.replace(statusRegex, '');
+
+  // Update the comment
+  comment.body = `[${status}] ${originalText}`;
+  comment.savedBody = originalText;
 
   // Add note comment about status change
+  const statusMessage = replyText ?
+    `Status changed to ${status}.\n\n${replyText}` :
+    `Status changed to ${status}.`;
+
   saveNoteComment(
     thread,
-    `Status changed to ${status}.`,
+    statusMessage,
     false,
     noteMap,
     author ? author : '',

src/models/noteComment.ts

@@ -9,7 +9,7 @@ export class NoteComment implements vscode.Comment {
   label: string | undefined;
   savedBody: string | vscode.MarkdownString; // for the Cancel button
   constructor(
-    public body: string,
+    public body: string | vscode.MarkdownString,
     public mode: vscode.CommentMode,
     public author: vscode.CommentAuthorInformation,
     public parent?: vscode.CommentThread,