fix(prompt-input): forward refs on the add menu and selects

PromptInputAddMenu and the model and mode selects read ref from their
props, which only React 19 passes to a function component; under
React 18, inside the peer range, the ref was dropped with a warning.
Wrap them, and the select they share, in forwardRef so a ref reaches
each part's root element on both versions.

Author: pratikbuilds

src/ui/prompt-input-toolbar.test.tsx

@@ -1,5 +1,5 @@
 import { describe, expect, spyOn, test } from "bun:test";
-import { act, type ReactNode, useState } from "react";
+import { act, createRef, type ReactNode, useState } from "react";
 import { createRoot, type Root } from "react-dom/client";
 
 import type { PromptToken } from "../lib/prompt-tokens.js";
@@ -613,6 +613,32 @@ describe("PromptInputButton", () => {
   });
 });
 
+describe("refs", () => {
+  test("the add menu and both selects hand their root to a ref", () => {
+    const add = createRef<HTMLDivElement>();
+    const model = createRef<HTMLDivElement>();
+    const mode = createRef<HTMLDivElement>();
+    const view = mount(
+      <>
+        <PromptInputAddMenu ref={add} items={menuItems([])} />
+        <PromptInputModelSelect ref={model} models={MODELS} />
+        <PromptInputModeSelect ref={mode} modes={MODES} />
+      </>,
+    );
+    expect(add.current).toBe(
+      view.query<HTMLDivElement>("[data-slot='prompt-input-add-menu']"),
+    );
+    expect(model.current).toBe(
+      view.query<HTMLDivElement>("[data-slot='prompt-input-model-select']"),
+    );
+    expect(mode.current).toBe(
+      view.query<HTMLDivElement>("[data-slot='prompt-input-mode-select']"),
+    );
+    view.unmount();
+    expect(add.current).toBeNull();
+  });
+});
+
 describe("PromptInputContext", () => {
   test("offers a clear button only when it can be cleared", () => {
     const cleared: string[] = [];

src/ui/prompt-input-toolbar.tsx

@@ -1,7 +1,14 @@
 "use client";
 
 import { Check, ChevronDown, Plus, Reply, X } from "lucide-react";
-import { useEffect, useId, useLayoutEffect, useRef, useState } from "react";
+import {
+  forwardRef,
+  useEffect,
+  useId,
+  useLayoutEffect,
+  useRef,
+  useState,
+} from "react";
 import type * as React from "react";
 
 import { useControllableState } from "../hooks/use-controllable-state.js";
@@ -234,18 +241,23 @@ export type PromptInputAddMenuProps = Omit<
  * elsewhere in the enclosing `PromptInput` (the textarea, beside a menu the
  * host opened).
  */
-export function PromptInputAddMenu({
-  items,
-  disabled: disabledProp,
-  label = "Add",
-  align = "start",
-  open: openProp,
-  defaultOpen = false,
-  onOpenChange,
-  className,
+export const PromptInputAddMenu = forwardRef<
+  HTMLDivElement,
+  PromptInputAddMenuProps
+>(function PromptInputAddMenu(
+  {
+    items,
+    disabled: disabledProp,
+    label = "Add",
+    align = "start",
+    open: openProp,
+    defaultOpen = false,
+    onOpenChange,
+    className,
+    ...props
+  },
   ref,
-  ...props
-}: PromptInputAddMenuProps) {
+) {
   const popover = useDismissablePopover<HTMLDivElement, HTMLButtonElement>({
     open: openProp,
     defaultOpen,
@@ -406,7 +418,7 @@ export function PromptInputAddMenu({
       ) : null}
     </div>
   );
-}
+});
 
 /** An option as the shared select draws it. */
 type SelectOption = {
@@ -474,22 +486,24 @@ type OptionSelectProps = SelectStateProps & {
  * `PromptInput` (the textarea, beside a list the host opened) closes it too
  * and leaves focus where it is.
  */
-function OptionSelect({
-  slot,
-  label,
-  options,
-  value: valueProp,
-  defaultValue,
-  onValueChange,
-  open: openProp,
-  defaultOpen = false,
-  onOpenChange,
-  disabled: disabledProp,
-  align = "start",
-  className,
-  ref,
-  ...props
-}: OptionSelectProps) {
+const OptionSelect = forwardRef(function OptionSelect(
+  {
+    slot,
+    label,
+    options,
+    value: valueProp,
+    defaultValue,
+    onValueChange,
+    open: openProp,
+    defaultOpen = false,
+    onOpenChange,
+    disabled: disabledProp,
+    align = "start",
+    className,
+    ...props
+  }: Omit<OptionSelectProps, "ref">,
+  ref: React.ForwardedRef<HTMLDivElement>,
+) {
   // Undefined until something is chosen: the first option then stands in,
   // read from the current list so options that load late still get it.
   const [picked, setPicked] = useControllableState<string | undefined>({
@@ -702,7 +716,7 @@ function OptionSelect({
       ) : null}
     </div>
   );
-}
+});
 
 /** A model to choose from in `PromptInputModelSelect`. */
 export type PromptInputModelOption = {
@@ -728,15 +742,17 @@ export type PromptInputModelSelectProps = SelectStateProps & {
  * reasoning level, is a separate control, so `onValueChange` stays a single
  * model id.
  */
-export function PromptInputModelSelect({
-  models,
-  label = "Model",
-  align = "end",
-  ...props
-}: PromptInputModelSelectProps) {
+export const PromptInputModelSelect = forwardRef<
+  HTMLDivElement,
+  PromptInputModelSelectProps
+>(function PromptInputModelSelect(
+  { models, label = "Model", align = "end", ...props },
+  ref,
+) {
   return (
     <OptionSelect
       {...props}
+      ref={ref}
       slot="prompt-input-model-select"
       label={label}
       align={align}
@@ -749,7 +765,7 @@ export function PromptInputModelSelect({
       }))}
     />
   );
-}
+});
 
 /** A way of working, offered by `PromptInputModeSelect`. */
 export type PromptInputModeOption = {
@@ -766,14 +782,14 @@ export type PromptInputModeSelectProps = SelectStateProps & {
 
 /** How the agent works this turn. The modes' labels and details are the
  * host's own strings; the trigger shows the chosen mode's label. */
-export function PromptInputModeSelect({
-  modes,
-  label = "Mode",
-  ...props
-}: PromptInputModeSelectProps) {
+export const PromptInputModeSelect = forwardRef<
+  HTMLDivElement,
+  PromptInputModeSelectProps
+>(function PromptInputModeSelect({ modes, label = "Mode", ...props }, ref) {
   return (
     <OptionSelect
       {...props}
+      ref={ref}
       slot="prompt-input-mode-select"
       label={label}
       options={modes.map((mode) => ({
@@ -784,7 +800,7 @@ export function PromptInputModeSelect({
       }))}
     />
   );
-}
+});
 
 export type PromptInputFloatProps = React.ComponentProps<"div">;