fix(web): return focus to the composer after saving a citation note

apps/web/src/components/ComposerPromptEditorTiptap.tsx

@@ -29,6 +29,7 @@ import {
   useMemo,
   useRef,
   useState,
+  type KeyboardEvent as ReactKeyboardEvent,
 } from "react";
 import { EditorContent, useEditor } from "@tiptap/react";
 
@@ -371,13 +372,40 @@ function ComposerCitationNodeView({ node, editor, getPos }: NodeViewProps) {
       .run();
   }, [editor, nodePos]);
 
+  // Put the caret right after the chip so Enter sends and typing continues the prompt.
+  const onRestoreFocus = useCallback(() => {
+    if (!editor.isEditable) return;
+    const pos = nodePos();
+    if (pos === null) return;
+    const current = editor.state.doc.nodeAt(pos);
+    if (!current || current.type.name !== "composer-citation") return;
+    editor.commands.focus(pos + current.nodeSize);
+  }, [editor, nodePos]);
+
   return (
     <NodeViewWrapper
       as="span"
       className="inline-flex min-w-0 max-w-full select-none"
       contentEditable={false}
       spellCheck={false}
       data-composer-citation-chip="true"
+      onKeyDown={(event: ReactKeyboardEvent<HTMLElement>) => {
+        // Tab from the comment button returns to the caret after the chip.
+        if (
+          !editor.isEditable ||
+          event.key !== "Tab" ||
+          event.shiftKey ||
+          event.altKey ||
+          event.metaKey ||
+          event.ctrlKey ||
+          !(event.target instanceof HTMLElement) ||
+          event.target.dataset.citationCommentTrigger === undefined
+        ) {
+          return;
+        }
+        event.preventDefault();
+        onRestoreFocus();
+      }}
     >
       <AssistantCitationChip
         citation={citation}
@@ -396,6 +424,7 @@ function ComposerCitationNodeView({ node, editor, getPos }: NodeViewProps) {
             commentContext.onSubmitAndSend();
             return true;
           },
+          onRestoreFocus,
         }}
       />
     </NodeViewWrapper>
@@ -856,6 +885,32 @@ function ComposerPromptEditorTiptapInner(props: ComposerPromptEditorProps) {
               return true;
             }
           }
+          // Shift+Tab from just after a citation reaches its comment button, which
+          // native tab order skips because the chip lives inside the editor.
+          if (
+            event.key === "Tab" &&
+            event.shiftKey &&
+            !event.altKey &&
+            !event.metaKey &&
+            !event.ctrlKey &&
+            view.state.selection.empty
+          ) {
+            const { $from } = view.state.selection;
+            const citation = $from.nodeBefore;
+            if (citation?.type.name === "composer-citation") {
+              const chip = view.nodeDOM($from.pos - citation.nodeSize);
+              const commentButton =
+                chip instanceof HTMLElement
+                  ? chip.querySelector<HTMLElement>("[data-citation-comment-trigger]")
+                  : null;
+              if (commentButton) {
+                event.preventDefault();
+                event.stopPropagation();
+                commentButton.focus();
+                return true;
+              }
+            }
+          }
           if (event.key === "Enter" && (event.isComposing || event.keyCode === 229)) {
             event.stopPropagation();
             return true;

apps/web/src/components/chat/AssistantCitationChip.tsx

@@ -39,10 +39,13 @@ export function AssistantCitationChip({
     onCancel?: () => void;
     onSave: (comment: string) => boolean;
     onSaveAndSend?: (comment: string) => boolean;
+    /** Returns focus to the host editor when the popover closes instead of to the pencil trigger. */
+    onRestoreFocus?: () => void;
   };
 }) {
   const navigate = useNavigate();
   const commentInputRef = useRef<HTMLTextAreaElement>(null);
+  const commentPopupRef = useRef<HTMLDivElement>(null);
   const draftCommentRef = useRef<string | null>(null);
   const [unavailableSourceAnchor, setUnavailableSourceAnchor] =
     useState<AssistantCitationSourceAnchor | null>(null);
@@ -155,6 +158,7 @@ export function AssistantCitationChip({
         >
           <PopoverTrigger
             aria-label={citation.comment ? "Edit citation comment" : "Add comment to citation"}
+            data-citation-comment-trigger="true"
             render={<ContextChipAction />}
           >
             <PencilIcon aria-hidden="true" />
@@ -169,6 +173,22 @@ export function AssistantCitationChip({
                 commentInputRef.current?.focus({ preventScroll: true });
                 return false;
               }}
+              finalFocus={
+                commentEditor.onRestoreFocus
+                  ? () => {
+                      // Leave focus alone when the user closed the popover by moving to another control.
+                      const activeElement = document.activeElement;
+                      if (
+                        activeElement === document.body ||
+                        (activeElement !== null && commentPopupRef.current?.contains(activeElement))
+                      ) {
+                        commentEditor.onRestoreFocus?.();
+                      }
+                      return false;
+                    }
+                  : undefined
+              }
+              ref={commentPopupRef}
               aria-label="Edit citation comment"
               width="md"
               padding="compact"