Ensure the "Reference in Prompt" command checks or partially checks directories when referencing them

Author: robertpiosik

packages/ui/src/components/editor/panel/ChatInput/hooks/use-handlers.ts

@@ -423,8 +423,11 @@ export const use_handlers = (
     }
 
     if (key == 'z' && (ctrlKey || metaKey) && !shiftKey) {
-      handle_undo_at_sign(e)
-      return
+      if (saved_value_before_at_sign !== null) {
+        handle_undo_at_sign(e)
+        return
+      }
+      // Otherwise, we don't return, allowing the default undo behavior
     }
 
     if (props.value) {

packages/vscode/src/commands/reference-in-prompt-command.ts

@@ -32,7 +32,10 @@ export const reference_in_prompt_command = (
         .get_all_checked_paths()
         .includes(file_path)
 
-      if (!is_checked) {
+      const is_partially_checked =
+        uri.isDirectory && workspace_provider.is_partially_checked(file_path)
+
+      if (!is_checked && !is_partially_checked) {
         await workspace_provider.update_check_state(
           uri,
           vscode.TreeItemCheckboxState.Checked

packages/vscode/src/context/providers/workspace-provider.ts

@@ -1328,6 +1328,10 @@ export class WorkspaceProvider
     return this.user_ignore.ignores(relative_path)
   }
 
+  public is_partially_checked(path: string): boolean {
+    return this.partially_checked_dirs.has(path)
+  }
+
   public async check_all(): Promise<void> {
     for (const workspace_root of this.workspace_roots) {
       this.checked_items.set(