fix(admin): keep keyboard focus in the wizard after adding a token inline

src/features/admin/components/wizard/sources/AddSourceFlow.tsx

@@ -6,6 +6,7 @@ import {
   useState,
   type CSSProperties,
   type ReactNode,
+  type Ref,
   type RefObject,
 } from "react";
 import { createPortal } from "react-dom";
@@ -22,6 +23,7 @@ import { JiraConnectStep } from "../../../../data-ingestion/components/JiraConne
 import { ConfluenceConnectStep } from "../../../../data-ingestion/components/ConfluenceConnectStep";
 import { SourceTypeStep } from "../../../../data-ingestion/components/SourceTypeStep";
 import { FileUploadZone } from "../../../../knowledge-base/components/FileUploadZone";
+import { useDialogFocus } from "../../../../../components/ui/useDialogFocus";
 import { TokenAddForm } from "../../../../settings/components/TokenAddForm";
 import { AtlassianCredentialAddForm } from "../../../../settings/components/atlassian/AtlassianCredentialAddForm";
 import type { SourceSystem } from "../../../../data-ingestion/types";
@@ -137,9 +139,23 @@ type AddSourceFlowProps = {
  * and closes itself once the token/credential is saved. The form owns save +
  * cancel; this only decides where it is shown.
  */
-function TriggerButton({ label, onClick }: { label: string; onClick: () => void }) {
+function TriggerButton({
+  label,
+  onClick,
+  ref,
+}: {
+  label: string;
+  onClick: () => void;
+  ref?: Ref<HTMLButtonElement>;
+}) {
   return (
-    <Button variant="secondary" size="sm" onClick={onClick} icon={<KeyRound className="h-4 w-4" />}>
+    <Button
+      ref={ref}
+      variant="secondary"
+      size="sm"
+      onClick={onClick}
+      icon={<KeyRound className="h-4 w-4" />}
+    >
       {label}
     </Button>
   );
@@ -171,6 +187,12 @@ const VIEWPORT_MARGIN = 12;
  * and it slides in/out. Escape closes only the companion — a capture-phase
  * handler stops the event before the wizard modal's own Escape-to-close fires.
  *
+ * Being portalled beside the wizard, it sits outside the wizard's Tab trap, so
+ * it runs its own: focus moves in on open, Tab stays inside while open, and
+ * focus returns to the trigger on close. Without that, a keyboard user who saved
+ * a token was left with focus on `<body>`, and the next Tab landed on the page
+ * underneath the wizard.
+ *
  * The position is measured from the wizard dialog (found via `anchorRef`) and
  * re-measured on scroll/resize; it clamps into the viewport if the modal sits
  * too far right for the full gap.
@@ -189,6 +211,7 @@ function CompanionModal({
   children: ReactNode;
 }) {
   const prefersReducedMotion = useReducedMotion();
+  const dialogRef = useDialogFocus<HTMLDivElement>(isOpen);
   const [style, setStyle] = useState<CSSProperties>({ position: "fixed", top: -9999, left: -9999 });
   // Remembers the last applied top/left so the per-frame tracker only triggers a
   // re-render when the position actually moves.
@@ -258,6 +281,8 @@ function CompanionModal({
     <AnimatePresence>
       {isOpen && (
         <motion.div
+          ref={dialogRef}
+          tabIndex={-1}
           role="dialog"
           aria-modal="false"
           aria-label={title}
@@ -322,8 +347,19 @@ function CredentialSlot({
   const isDesktop = useMediaQuery(DESKTOP_QUERY);
   const [isOpen, setIsOpen] = useState(false);
   const anchorRef = useRef<HTMLDivElement>(null);
+  const triggerRef = useRef<HTMLButtonElement>(null);
+  const wasOpenRef = useRef(false);
   const close = () => setIsOpen(false);
 
+  // Hand focus back to the trigger once the form closes (saved or cancelled).
+  // The focused control was inside the form, so without this focus drops to
+  // `<body>` and the next Tab escapes the wizard onto the page behind it. On
+  // phone the trigger remounts in place of the inline form, hence the effect.
+  useEffect(() => {
+    if (wasOpenRef.current && !isOpen) triggerRef.current?.focus();
+    wasOpenRef.current = isOpen;
+  }, [isOpen]);
+
   const companionOpen = isDesktop && isOpen;
   useEffect(() => {
     onCompanionOpenChange?.(companionOpen);
@@ -337,15 +373,15 @@ function CredentialSlot({
       <>{renderForm(close, false)}</>
     ) : (
       <div className="flex flex-wrap items-center gap-3">
-        <TriggerButton label={buttonLabel} onClick={() => setIsOpen(true)} />
+        <TriggerButton ref={triggerRef} label={buttonLabel} onClick={() => setIsOpen(true)} />
         {missingLabel && <MissingCredentialChip label={missingLabel} />}
       </div>
     );
   }
 
   return (
     <div ref={anchorRef} className="flex flex-wrap items-center gap-3">
-      <TriggerButton label={buttonLabel} onClick={() => setIsOpen(true)} />
+      <TriggerButton ref={triggerRef} label={buttonLabel} onClick={() => setIsOpen(true)} />
       {missingLabel && <MissingCredentialChip label={missingLabel} />}
       <CompanionModal isOpen={isOpen} onClose={close} title={panelTitle} anchorRef={anchorRef}>
         {renderForm(close, true)}

tests/unit/features/admin/components/CreateProjectWizard.test.tsx

@@ -1,4 +1,4 @@
-import { describe, it, expect, vi, beforeEach } from "vitest";
+import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
 import { act, render, screen, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { CreateProjectWizard } from "../../../../../src/features/admin/components/CreateProjectWizard";
@@ -767,6 +767,81 @@ describe("CreateProjectWizard", () => {
     );
   });
 
+  it("returns focus to the trigger after a token is added inline by keyboard", async () => {
+    vi.mocked(getGithubPatNames).mockResolvedValueOnce([]).mockResolvedValue(["fresh-pat"]);
+    const user = userEvent.setup();
+    renderWizard({ tokenNames: [] });
+
+    await goToSources(user);
+    await openGithubDetail(user);
+
+    screen.getByRole("button", { name: /add github token/i }).focus();
+    await user.keyboard("{Enter}");
+    await user.type(screen.getByTestId("settings-add-token-name"), "fresh-pat");
+    await user.type(screen.getByTestId("settings-add-token-value"), "ghp_secret123{Enter}");
+
+    // Focus lands back on the trigger instead of dropping to <body>, where the
+    // next Tab would escape the wizard onto the page behind it.
+    await waitFor(() =>
+      expect(screen.getByRole("button", { name: /add github token/i })).toHaveFocus(),
+    );
+  });
+
+  describe("on desktop, where the token form opens in a companion beside the wizard", () => {
+    beforeEach(() => {
+      vi.spyOn(window, "matchMedia").mockImplementation(
+        (query: string) =>
+          ({
+            matches: query === "(min-width: 1280px)",
+            media: query,
+            onchange: null,
+            addEventListener: vi.fn(),
+            removeEventListener: vi.fn(),
+            addListener: vi.fn(),
+            removeListener: vi.fn(),
+            dispatchEvent: vi.fn(),
+          }) as MediaQueryList,
+      );
+    });
+
+    afterEach(() => {
+      vi.restoreAllMocks();
+    });
+
+    it("keeps Tab inside the companion and returns focus to the trigger after saving", async () => {
+      vi.mocked(getGithubPatNames).mockResolvedValueOnce([]).mockResolvedValue(["fresh-pat"]);
+      const user = userEvent.setup();
+      renderWizard({ tokenNames: [] });
+
+      await goToSources(user);
+      await openGithubDetail(user);
+
+      const trigger = screen.getByRole("button", { name: /add github token/i });
+      trigger.focus();
+      await user.keyboard("{Enter}");
+
+      const companion = await screen.findByRole("dialog", { name: "New GitHub token" });
+      await settleModalFocus();
+      expect(screen.getByTestId("settings-add-token-name")).toHaveFocus();
+
+      // Tabbing past the last control wraps within the companion.
+      const submit = screen.getByTestId("settings-add-token-submit");
+      await user.type(screen.getByTestId("settings-add-token-name"), "fresh-pat");
+      await user.type(screen.getByTestId("settings-add-token-value"), "ghp_secret123");
+      submit.focus();
+      await user.tab();
+      expect(companion).toContainElement(document.activeElement as HTMLElement);
+
+      submit.focus();
+      await user.keyboard("{Enter}");
+
+      await waitFor(() =>
+        expect(vi.mocked(addGithubPat)).toHaveBeenCalledWith("fresh-pat", "ghp_secret123"),
+      );
+      await waitFor(() => expect(trigger).toHaveFocus());
+    });
+  });
+
   it("adds a Jira credential inline and selects the new one", async () => {
     vi.mocked(getMyAtlassianCredentials)
       .mockResolvedValueOnce([]) // initial load: none stored