fix(ui): improve picker actions, directory navigation, @ handling, and message display

packages/ui/src/components/prompt-input/usePromptKeyDown.ts

@@ -183,9 +183,25 @@ export function usePromptKeyDown(options: UsePromptKeyDownOptions) {
 
         if (isDeletingFromEnd || isDeletingFromStart || isSelected) {
           const currentAttachments = options.getAttachments()
-          const attachment = currentAttachments.find(
-            (a) => (a.source.type === "file" || a.source.type === "agent") && a.filename === name,
-          )
+          const attachment = currentAttachments.find((a) => {
+            if (a.source.type === "agent") {
+              return a.filename === name
+            }
+            if (a.source.type === "file") {
+              // Match either by filename (basename) or by path (for full paths like @docs/file.txt)
+              return (
+                a.filename === name ||
+                a.source.path === name ||
+                a.source.path.endsWith("/" + name) ||
+                a.source.path === name.replace(/\/$/, "")
+              )
+            }
+            if (a.source.type === "text") {
+              // For text attachments (path-only mentions), match by value
+              return a.source.value === name || a.source.value.endsWith("/" + name)
+            }
+            return false
+          })
 
           if (attachment) {
             e.preventDefault()
@@ -205,6 +221,14 @@ export function usePromptKeyDown(options: UsePromptKeyDownOptions) {
               textarea.setSelectionRange(mentionStart, mentionStart)
             }, 0)
 
+            // Check if there are any @ remaining in the text - if not, close the picker
+            if (!newText.includes("@") && options.isPickerOpen()) {
+              options.closePicker()
+              // Clear ignoredAtPositions since we deleted the entire @mention
+              // This ensures typing @ again will open the picker
+              options.setIgnoredAtPositions(new Set())
+            }
+
             return
           }
         }

packages/ui/src/components/prompt-input/usePromptPicker.ts

@@ -236,7 +236,7 @@ export function usePromptPicker(options: PromptPickerOptions): PromptPickerContr
         const mentionText = `@${folderMention}`
 
         if (action === "shiftEnter") {
-          // SHIFT+ENTER on directory: attach path as text only.
+          // SHIFT+ENTER on directory: keep @path in prompt, add text attachment, remove @ when sending
           addPathOnlyAttachment(folderMention)
           replaceMentionToken(mentionText, { trailingSpace: true })
         } else {
@@ -274,7 +274,7 @@ export function usePromptPicker(options: PromptPickerOptions): PromptPickerContr
         }
 
         if (action === "shiftEnter") {
-          // SHIFT+ENTER on file: attach path as text only.
+          // SHIFT+ENTER on file: keep @path in prompt, add text attachment, remove @ when sending
           addPathOnlyAttachment(normalizedPath)
           replaceMentionToken(`@${normalizedPath}`, { trailingSpace: true })
         } else {
@@ -316,6 +316,28 @@ export function usePromptPicker(options: PromptPickerOptions): PromptPickerContr
     const pos = atPosition()
     if (pickerMode() === "mention" && pos !== null) {
       setIgnoredAtPositions((prev) => new Set(prev).add(pos))
+
+      // Remove the partial @mention text from the textarea when ESC is pressed
+      const textarea = options.getTextarea()
+      if (textarea) {
+        const currentPrompt = options.prompt()
+        const cursorPos = textarea.selectionStart
+        // Remove text from @ position to cursor position
+        const before = currentPrompt.substring(0, pos)
+        const after = currentPrompt.substring(cursorPos)
+        options.setPrompt(before + after)
+
+        // Restore cursor position to where @ was
+        setTimeout(() => {
+          const nextTextarea = options.getTextarea()
+          if (nextTextarea) {
+            nextTextarea.setSelectionRange(pos, pos)
+          }
+        }, 0)
+
+        // Clear ignoredAtPositions so typing @ again will work
+        setIgnoredAtPositions(new Set<number>())
+      }
     }
     setShowPicker(false)
     setAtPosition(null)