From 0c4ec2d3747b74bb96377e69275e896b7203f1a7 Mon Sep 17 00:00:00 2001
From: MacKinley Smith <smithmackinley@gmail.com>
Date: Thu, 24 Sep 2026 10:32:24 -0600
Subject: [PATCH 1/2] fix(web): return focus to the composer after saving a
 citation note

Saving or cancelling a citation comment now puts the caret right after the chip instead of on the edit button, so Enter sends. Shift+Tab from just after a chip reaches its edit button, and Tab from that button returns to the caret.
---
 .../components/ComposerPromptEditorTiptap.tsx | 54 +++++++++++++++++++
 .../components/chat/AssistantCitationChip.tsx | 20 ++++++-
 2 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/apps/web/src/components/ComposerPromptEditorTiptap.tsx b/apps/web/src/components/ComposerPromptEditorTiptap.tsx
index 50305ffdbef0..acd2220c25c1 100644
--- a/apps/web/src/components/ComposerPromptEditorTiptap.tsx
+++ b/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,6 +372,16 @@ 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"
@@ -378,6 +389,22 @@ function ComposerCitationNodeView({ node, editor, getPos }: NodeViewProps) {
       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 (
+          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 +423,7 @@ function ComposerCitationNodeView({ node, editor, getPos }: NodeViewProps) {
             commentContext.onSubmitAndSend();
             return true;
           },
+          onRestoreFocus,
         }}
       />
     </NodeViewWrapper>
@@ -856,6 +884,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;
diff --git a/apps/web/src/components/chat/AssistantCitationChip.tsx b/apps/web/src/components/chat/AssistantCitationChip.tsx
index d838c53d27d2..8558f29f36c0 100644
--- a/apps/web/src/components/chat/AssistantCitationChip.tsx
+++ b/apps/web/src/components/chat/AssistantCitationChip.tsx
@@ -23,7 +23,7 @@ import { Popover, PopoverPopup, PopoverTrigger } from "../ui/popover";
 import { AssistantCitationCommentEditor } from "./AssistantCitationCommentEditor";
 import { resolveAssistantCitationCommentDismissal } from "./assistantCitationCommentDismissal";
 import { observeAssistantCitationCommentSource } from "./AssistantCitationSource";
-import { composerFloatingLayerProps } from "./composerEventScope";
+import { composerFloatingLayerProps, isInsideComposerFloatingLayer } from "./composerEventScope";
 
 export function AssistantCitationChip({
   citation,
@@ -39,6 +39,8 @@ 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();
@@ -155,6 +157,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 +172,21 @@ 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 ||
+                        isInsideComposerFloatingLayer(activeElement)
+                      ) {
+                        commentEditor.onRestoreFocus?.();
+                      }
+                      return false;
+                    }
+                  : undefined
+              }
               aria-label="Edit citation comment"
               width="md"
               padding="compact"

From eb544d2756817a22a23b68bd768b4ee8e5c20123 Mon Sep 17 00:00:00 2001
From: MacKinley Smith <smithmackinley@gmail.com>
Date: Thu, 24 Sep 2026 10:53:08 -0600
Subject: [PATCH 2/2] fix(web): only restore focus from the citation popover
 itself

Restoring focus whenever it sat in any composer floating layer could pull it out of another composer menu. Also leave Tab alone on a read-only composer instead of swallowing it.
---
 apps/web/src/components/ComposerPromptEditorTiptap.tsx | 1 +
 apps/web/src/components/chat/AssistantCitationChip.tsx | 6 ++++--
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/apps/web/src/components/ComposerPromptEditorTiptap.tsx b/apps/web/src/components/ComposerPromptEditorTiptap.tsx
index acd2220c25c1..678f250bcab3 100644
--- a/apps/web/src/components/ComposerPromptEditorTiptap.tsx
+++ b/apps/web/src/components/ComposerPromptEditorTiptap.tsx
@@ -392,6 +392,7 @@ function ComposerCitationNodeView({ node, editor, getPos }: NodeViewProps) {
       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 ||
diff --git a/apps/web/src/components/chat/AssistantCitationChip.tsx b/apps/web/src/components/chat/AssistantCitationChip.tsx
index 8558f29f36c0..c6a9ce443938 100644
--- a/apps/web/src/components/chat/AssistantCitationChip.tsx
+++ b/apps/web/src/components/chat/AssistantCitationChip.tsx
@@ -23,7 +23,7 @@ import { Popover, PopoverPopup, PopoverTrigger } from "../ui/popover";
 import { AssistantCitationCommentEditor } from "./AssistantCitationCommentEditor";
 import { resolveAssistantCitationCommentDismissal } from "./assistantCitationCommentDismissal";
 import { observeAssistantCitationCommentSource } from "./AssistantCitationSource";
-import { composerFloatingLayerProps, isInsideComposerFloatingLayer } from "./composerEventScope";
+import { composerFloatingLayerProps } from "./composerEventScope";
 
 export function AssistantCitationChip({
   citation,
@@ -45,6 +45,7 @@ export function AssistantCitationChip({
 }) {
   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);
@@ -179,7 +180,7 @@ export function AssistantCitationChip({
                       const activeElement = document.activeElement;
                       if (
                         activeElement === document.body ||
-                        isInsideComposerFloatingLayer(activeElement)
+                        (activeElement !== null && commentPopupRef.current?.contains(activeElement))
                       ) {
                         commentEditor.onRestoreFocus?.();
                       }
@@ -187,6 +188,7 @@ export function AssistantCitationChip({
                     }
                   : undefined
               }
+              ref={commentPopupRef}
               aria-label="Edit citation comment"
               width="md"
               padding="compact"