From 11b444a3ec30bf7eca233991490fe38e506505d8 Mon Sep 17 00:00:00 2001
From: Peter Vachon <peteraldenvachon@gmail.com>
Date: Thu, 30 Apr 2026 08:18:26 -0400
Subject: [PATCH] feat(apollo-vertex): ai-chat text selection Ask AI [5/5]

Adds text-selection-to-AI flow:
- AiChatSelectionMenu: floating "Ask AI" pill shown on text selection,
  scoped to the chat container, with fade+scale entrance animation
- AiChatInput: quotedText chip with left-accent border, dismissable with X
- AiChat: enableTextSelection prop wires selection menu to quote chip;
  clicking "Ask AI" pre-fills the input and focuses the textarea
- registry.json: ai-chat-selection-menu.tsx registered

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 apps/apollo-vertex/registry.json              |   5 +
 .../ai-chat/components/ai-chat-input.tsx      |  26 +++-
 .../components/ai-chat-selection-menu.tsx     | 112 ++++++++++++++++++
 .../registry/ai-chat/components/ai-chat.tsx   |  32 ++++-
 .../templates/ai-chat/AiChatAgentHubMode.tsx  |  25 +++-
 5 files changed, 195 insertions(+), 5 deletions(-)
 create mode 100644 apps/apollo-vertex/registry/ai-chat/components/ai-chat-selection-menu.tsx

diff --git a/apps/apollo-vertex/registry.json b/apps/apollo-vertex/registry.json
index b5d85c435..baf2690b8 100644
--- a/apps/apollo-vertex/registry.json
+++ b/apps/apollo-vertex/registry.json
@@ -443,6 +443,11 @@
           "type": "registry:ui",
           "target": "components/ui/ai-chat/components/ai-chat-message-actions.tsx"
         },
+        {
+          "path": "registry/ai-chat/components/ai-chat-selection-menu.tsx",
+          "type": "registry:ui",
+          "target": "components/ui/ai-chat/components/ai-chat-selection-menu.tsx"
+        },
         {
           "path": "registry/ai-chat/components/ai-chat-thinking.tsx",
           "type": "registry:ui",
diff --git a/apps/apollo-vertex/registry/ai-chat/components/ai-chat-input.tsx b/apps/apollo-vertex/registry/ai-chat/components/ai-chat-input.tsx
index d8df47b38..45c93514a 100644
--- a/apps/apollo-vertex/registry/ai-chat/components/ai-chat-input.tsx
+++ b/apps/apollo-vertex/registry/ai-chat/components/ai-chat-input.tsx
@@ -8,12 +8,12 @@ import {
   type FormEvent,
   type KeyboardEvent,
   type Ref,
-  createPortal,
   useEffect,
   useImperativeHandle,
   useRef,
   useState,
 } from "react";
+import { createPortal } from "react-dom";
 import { useTranslation } from "react-i18next";
 import {
   DropdownMenu,
@@ -65,6 +65,10 @@ interface AiChatInputProps {
   placeholder?: string;
   hasMessages?: boolean;
   maxLength?: number;
+  /** Pre-filled quoted text shown as a chip above the input */
+  quotedText?: string;
+  /** Called when the user dismisses the quoted text chip */
+  onClearQuote?: () => void;
   ref?: Ref<AiChatInputHandle>;
 }
 
@@ -82,6 +86,8 @@ export function AiChatInput({
   placeholder,
   hasMessages = false,
   maxLength,
+  quotedText,
+  onClearQuote,
   ref,
 }: AiChatInputProps) {
   const { t } = useTranslation();
@@ -246,6 +252,22 @@ export function AiChatInput({
       </div>
     ) : null;
 
+  const quoteChip = quotedText ? (
+    <div className="flex items-center px-3 pt-2">
+      <div className="inline-flex items-center gap-1.5 rounded-lg bg-ai-chat-muted px-2 py-1.5 text-xs text-ai-chat-muted-foreground max-w-full min-w-0">
+        <span className="truncate">{quotedText}</span>
+        <button
+          type="button"
+          onClick={onClearQuote}
+          className="flex-shrink-0 rounded hover:bg-ai-chat-border transition-colors p-0.5"
+          aria-label="Remove quoted text"
+        >
+          <X className="size-3" aria-hidden="true" />
+        </button>
+      </div>
+    </div>
+  ) : null;
+
   const plusMenu = (
     <DropdownMenu>
       <DropdownMenuTrigger asChild>
@@ -341,6 +363,7 @@ export function AiChatInput({
         </div>
         {hasMessages ? (
           <form {...formProps}>
+            {quoteChip}
             {fileChips}
             <div className="flex items-end gap-2 pl-[8px] pr-[8px] pt-[4px] pb-[8px]">
               {plusMenu}
@@ -362,6 +385,7 @@ export function AiChatInput({
           </form>
         ) : (
           <form {...formProps}>
+            {quoteChip}
             {fileChips}
             <textarea
               ref={textareaRef}
diff --git a/apps/apollo-vertex/registry/ai-chat/components/ai-chat-selection-menu.tsx b/apps/apollo-vertex/registry/ai-chat/components/ai-chat-selection-menu.tsx
new file mode 100644
index 000000000..14de1ca46
--- /dev/null
+++ b/apps/apollo-vertex/registry/ai-chat/components/ai-chat-selection-menu.tsx
@@ -0,0 +1,112 @@
+"use client";
+
+import { motion } from "framer-motion";
+import { useCallback, useEffect, useRef, useState } from "react";
+import { AutopilotIcon } from "./icons/autopilot";
+
+interface SelectionPosition {
+  x: number;
+  y: number;
+}
+
+interface AiChatSelectionMenuProps {
+  /** Called when the user clicks "Ask AI" — receives the selected text */
+  onAskAi: (selectedText: string) => void;
+  /** Container element to scope selection detection to, defaults to document */
+  containerRef?: React.RefObject<HTMLElement | null>;
+}
+
+export function AiChatSelectionMenu({
+  onAskAi,
+  containerRef,
+}: AiChatSelectionMenuProps) {
+  const [position, setPosition] = useState<SelectionPosition | null>(null);
+  const [selectedText, setSelectedText] = useState("");
+  const menuRef = useRef<HTMLDivElement>(null);
+
+  const handleSelectionChange = useCallback(() => {
+    const selection = window.getSelection();
+    if (!selection || selection.isCollapsed || !selection.toString().trim()) {
+      setPosition(null);
+      setSelectedText("");
+      return;
+    }
+
+    const text = selection.toString().trim();
+    const range = selection.getRangeAt(0);
+    const rect = range.getBoundingClientRect();
+
+    // If a container is provided, only show the menu for selections inside it
+    if (containerRef?.current) {
+      const container = containerRef.current;
+      const node = range.commonAncestorContainer;
+      if (!container.contains(node)) {
+        setPosition(null);
+        setSelectedText("");
+        return;
+      }
+    }
+
+    setSelectedText(text);
+    setPosition({
+      // Center the menu above the selection, clamped to viewport
+      x: Math.min(
+        Math.max(rect.left + rect.width / 2, 80),
+        window.innerWidth - 80,
+      ),
+      y: rect.top - 8,
+    });
+  }, [containerRef]);
+
+  useEffect(() => {
+    document.addEventListener("selectionchange", handleSelectionChange);
+    return () =>
+      document.removeEventListener("selectionchange", handleSelectionChange);
+  }, [handleSelectionChange]);
+
+  // Dismiss on click outside the menu
+  useEffect(() => {
+    if (!position) return;
+    const handlePointerDown = (e: PointerEvent) => {
+      if (menuRef.current && !menuRef.current.contains(e.target as Node)) {
+        setPosition(null);
+        setSelectedText("");
+      }
+    };
+    document.addEventListener("pointerdown", handlePointerDown);
+    return () => document.removeEventListener("pointerdown", handlePointerDown);
+  }, [position]);
+
+  if (!position) return null;
+
+  return (
+    <motion.div
+      ref={menuRef}
+      initial={{ opacity: 0, y: 4, scale: 0.96 }}
+      animate={{ opacity: 1, y: 0, scale: 1 }}
+      exit={{ opacity: 0, y: 4, scale: 0.96 }}
+      transition={{ duration: 0.14, ease: [0.22, 1, 0.36, 1] }}
+      className="fixed z-50 -translate-x-1/2 -translate-y-full"
+      style={{ left: position.x, top: position.y }}
+    >
+      <button
+        type="button"
+        onPointerDown={(e) => {
+          // Prevent the selection from clearing before we capture the text
+          e.preventDefault();
+        }}
+        onClick={() => {
+          onAskAi(selectedText);
+          window.getSelection()?.removeAllRanges();
+          setPosition(null);
+          setSelectedText("");
+        }}
+        className="flex items-center gap-1.5 px-3 py-1.5 rounded-full text-xs font-semibold text-white shadow-lg transition-opacity hover:opacity-90"
+        style={{ background: "var(--ai-gradient-strong)" }}
+      >
+        <AutopilotIcon size={14} aria-hidden="true" />
+        {"Ask AI"}
+      </button>
+    </motion.div>
+  );
+}
diff --git a/apps/apollo-vertex/registry/ai-chat/components/ai-chat.tsx b/apps/apollo-vertex/registry/ai-chat/components/ai-chat.tsx
index 91d71e495..edc3f516f 100644
--- a/apps/apollo-vertex/registry/ai-chat/components/ai-chat.tsx
+++ b/apps/apollo-vertex/registry/ai-chat/components/ai-chat.tsx
@@ -7,7 +7,13 @@ import {
   MoreHorizontal,
   RefreshCw,
 } from "lucide-react";
-import { type ReactNode, useEffect, useRef, useState } from "react";
+import {
+  type ReactNode,
+  useCallback,
+  useEffect,
+  useRef,
+  useState,
+} from "react";
 import { useTranslation } from "react-i18next";
 import {
   AlertDialog,
@@ -31,6 +37,7 @@ import type { MessageFeedbackType } from "../types";
 import { AiChatInput, type AiChatInputHandle } from "./ai-chat-input";
 import { AiChatLoading } from "./ai-chat-loading";
 import { AiChatProvider } from "./ai-chat-provider";
+import { AiChatSelectionMenu } from "./ai-chat-selection-menu";
 import { AutopilotGradientIcon } from "./icons/autopilot-gradient";
 
 export interface AiChatProps {
@@ -59,6 +66,8 @@ export interface AiChatProps {
   showClearButton?: boolean;
   showMessageActions?: boolean;
   showCopyButton?: boolean;
+  /** Show an "Ask AI" popup button when the user selects text inside the chat */
+  enableTextSelection?: boolean;
   error?: Error | null;
 }
 
@@ -83,12 +92,15 @@ export function AiChat({
   showClearButton = true,
   showMessageActions = true,
   showCopyButton = true,
+  enableTextSelection = false,
   error,
 }: AiChatProps) {
   const { t } = useTranslation();
   const [input, setInput] = useState("");
+  const [quotedText, setQuotedText] = useState<string | undefined>(undefined);
   const { scrollRef, contentRef, isStuck, scrollToBottom } = useStickyScroll();
   const inputRef = useRef<AiChatInputHandle>(null);
+  const chatContainerRef = useRef<HTMLDivElement>(null);
 
   const displayName = assistantName ?? t("ai_assistant");
 
@@ -137,6 +149,7 @@ export function AiChat({
     }
     onSendMessage(input.trim(), files);
     setInput("");
+    setQuotedText(undefined);
     scrollToBottom();
   };
 
@@ -155,6 +168,11 @@ export function AiChat({
     wasLoadingRef.current = isLoading;
   }, [isLoading, onSendMessage, scrollToBottom]);
 
+  const handleAskAi = useCallback((text: string) => {
+    setQuotedText(text);
+    requestAnimationFrame(() => inputRef.current?.focus());
+  }, []);
+
   const latestAssistantMessageId =
     messages.findLast((m) => m.role === "assistant")?.id ?? null;
 
@@ -195,9 +213,17 @@ export function AiChat({
       onEditMessage={onEditMessage}
     >
       <div
+        ref={chatContainerRef}
         className="flex flex-col h-full max-w-[680px] mx-auto bg-transparent text-ai-chat-foreground overflow-hidden"
         data-slot="ai-chat"
       >
+        {enableTextSelection && (
+          <AiChatSelectionMenu
+            onAskAi={handleAskAi}
+            containerRef={chatContainerRef}
+          />
+        )}
+
         {header ??
           (title && (
             <div className="relative z-10 py-3 px-4 flex items-center justify-between gap-2 bg-background">
@@ -279,6 +305,8 @@ export function AiChat({
                 isLoading={isLoading}
                 placeholder={placeholder}
                 hasMessages={false}
+                quotedText={quotedText}
+                onClearQuote={() => setQuotedText(undefined)}
               />
               {suggestions && suggestions.length > 0 && (
                 <div className="mt-4 px-4 flex flex-wrap justify-center gap-2">
@@ -374,6 +402,8 @@ export function AiChat({
               isLoading={isLoading}
               placeholder={placeholder}
               hasMessages
+              quotedText={quotedText}
+              onClearQuote={() => setQuotedText(undefined)}
             />
             <div className="pt-2 pb-3 px-4 text-xs leading-normal text-muted-foreground text-center">
               {t("ai_response_disclaimer")}
diff --git a/apps/apollo-vertex/templates/ai-chat/AiChatAgentHubMode.tsx b/apps/apollo-vertex/templates/ai-chat/AiChatAgentHubMode.tsx
index b3c9f6a5b..ffe2b0e3b 100644
--- a/apps/apollo-vertex/templates/ai-chat/AiChatAgentHubMode.tsx
+++ b/apps/apollo-vertex/templates/ai-chat/AiChatAgentHubMode.tsx
@@ -127,11 +127,27 @@ function AgentHubChatInner({
     tools,
   });
 
-  const { messages, sendMessage, isLoading, stop, clear, error } = useChat({
+  const {
+    messages,
+    sendMessage,
+    isLoading,
+    stop,
+    clear,
+    reload,
+    setMessages,
+    error,
+  } = useChat({
     connection,
     tools,
   });
 
+  const handleEditMessage = (messageId: string, content: string) => {
+    const idx = messages.findIndex((m) => m.id === messageId);
+    if (idx === -1) return;
+    setMessages(messages.slice(0, idx));
+    void sendMessage(content);
+  };
+
   const emptyState = (
     <AiChatEmptyState
       description={t("autopilot_empty_description")}
@@ -148,8 +164,11 @@ function AgentHubChatInner({
       }}
       onStop={stop}
       onClearChat={clear}
-      title="Autopilot"
-      assistantName="Autopilot"
+      onRegenerate={() => void reload()}
+      onEditMessage={handleEditMessage}
+      enableTextSelection
+      title="AI Assistant"
+      assistantName="AI Assistant"
       emptyState={emptyState}
       suggestions={[
         t("shell_suggestion_recent_runs"),