fix(web): keep the chat composer focused after sending a message (#3537)

Co-authored-by: nova28 <17953305+nova28@users.noreply.github.com>
Co-authored-by: Carlos Florencio <carlosmflorencio@gmail.com>

apps/web/components/task/chat/chat-input-toolbar-primitives.tsx

@@ -103,6 +103,7 @@ function SendSubmitButton({
             planModeEnabled && "bg-violet-600 hover:bg-violet-500",
           )}
           disabled={isDisabled}
+          onMouseDown={(e) => e.preventDefault()}
           onClick={onSubmit}
           data-testid="submit-message-button"
         >

apps/web/components/task/chat/use-tiptap-editor.test.ts

@@ -1,9 +1,12 @@
-import { describe, expect, it } from "vitest";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { renderHook } from "@testing-library/react";
 import { Extension } from "@tiptap/core";
 import {
   TIPTAP_EDITOR_TEXT_SIZE_CLASS,
   buildEditorExtensions,
   decideSubmitShortcut,
+  shouldRestoreFocusOnEnable,
+  useSyncDisabledState,
 } from "./use-tiptap-editor";
 import * as tiptapEditor from "./use-tiptap-editor";
 import { decideHistoryNav } from "./tiptap-editor-history";
@@ -166,6 +169,269 @@ describe("decideSubmitShortcut", () => {
   });
 });
 
+// Regression: after a send, ProseMirror flips `contenteditable` false then
+// true. A real browser blurs on the first flip and does not restore focus on
+// the second (jsdom does not reproduce this, so the blur is staged
+// explicitly here) -- the composer must regain focus unless something else
+// has since claimed it.
+describe("shouldRestoreFocusOnEnable", () => {
+  let input: HTMLInputElement;
+  let other: HTMLInputElement;
+
+  afterEach(() => {
+    input.remove();
+    other.remove();
+  });
+
+  function mountInputs() {
+    input = document.createElement("input");
+    other = document.createElement("input");
+    document.body.append(input, other);
+  }
+
+  it("restores focus when nothing has since claimed it", () => {
+    mountInputs();
+    input.focus();
+    expect(document.activeElement).toBe(input);
+    input.blur();
+    expect(document.activeElement).toBe(document.body);
+
+    expect(shouldRestoreFocusOnEnable(true)).toBe(true);
+  });
+
+  it("does not restore focus once another element has claimed it", () => {
+    mountInputs();
+    input.focus();
+    input.blur();
+    other.focus();
+    expect(document.activeElement).toBe(other);
+
+    expect(shouldRestoreFocusOnEnable(true)).toBe(false);
+  });
+
+  it("does nothing when the editor never had focus before disabling", () => {
+    mountInputs();
+    input.blur();
+    expect(document.activeElement).toBe(document.body);
+
+    expect(shouldRestoreFocusOnEnable(false)).toBe(false);
+  });
+});
+
+// Regression: exercises the actual effect wiring, not just the pure
+// predicate above -- a mock editor stands in for TipTap's `Editor` since
+// jsdom cannot reproduce the browser's disable-blurs-the-element behavior
+// the effect exists to work around.
+//
+// requestAnimationFrame is stubbed with a queue the test drives one frame at
+// a time (rather than firing synchronously), so a test can interleave a
+// browser-driven blur landing *between* retry attempts -- the scenario that
+// made the real fix necessary (a same-tick restore raced a queued blur and
+// lost).
+let frameQueue: FrameRequestCallback[];
+
+function stubAnimationFrame() {
+  frameQueue = [];
+  vi.stubGlobal("requestAnimationFrame", (cb: FrameRequestCallback) => {
+    frameQueue.push(cb);
+    return frameQueue.length;
+  });
+  vi.stubGlobal("cancelAnimationFrame", (handle: number) => {
+    frameQueue[handle - 1] = () => {};
+  });
+}
+
+function flushFrame() {
+  const cb = frameQueue.shift();
+  cb?.(0);
+}
+
+/** Stateful stand-in for TipTap's `Editor`: `view.hasFocus()` reflects
+ *  whatever last happened to it, `setEditable(false)` simulates the real
+ *  browser's disable-blurs-the-element behavior, and `commands.focus`
+ *  simulates a successful restore -- so a test can also simulate a
+ *  *delayed* external blur landing after a restore already ran. */
+function makeEditor(hasFocusInitially: boolean) {
+  let focused = hasFocusInitially;
+  return {
+    view: { hasFocus: () => focused },
+    setEditable: vi.fn((editable: boolean) => {
+      if (!editable) focused = false;
+    }),
+    commands: { focus: vi.fn(() => void (focused = true)) },
+    simulateExternalBlur: () => void (focused = false),
+  };
+}
+
+/** Variant of `makeEditor` where `setEditable(false)` does NOT blur --
+ *  modeling the real Chromium behavior the retry loop's comment cites, where
+ *  the blur caused by `contenteditable` flipping is applied on a queued task
+ *  rather than synchronously. `landQueuedBlur()` fires that queued blur
+ *  on demand, independent of `setEditable`, so a test can land it after the
+ *  loop has already run one or more attempts. */
+function makeQueuedBlurEditor(hasFocusInitially: boolean) {
+  let focused = hasFocusInitially;
+  return {
+    view: { hasFocus: () => focused },
+    setEditable: vi.fn(),
+    commands: { focus: vi.fn(() => void (focused = true)) },
+    landQueuedBlur: () => void (focused = false),
+  };
+}
+
+describe("useSyncDisabledState", () => {
+  afterEach(() => {
+    vi.unstubAllGlobals();
+  });
+
+  it("captures focus before disabling, then restores it on re-enable", () => {
+    stubAnimationFrame();
+    const editor = makeEditor(true);
+    const { rerender } = renderHook(({ disabled }) => useSyncDisabledState(editor, disabled), {
+      initialProps: { disabled: false },
+    });
+
+    rerender({ disabled: true });
+    expect(editor.setEditable).toHaveBeenLastCalledWith(false);
+    expect(editor.commands.focus).not.toHaveBeenCalled();
+
+    rerender({ disabled: false });
+    expect(editor.setEditable).toHaveBeenLastCalledWith(true);
+    expect(editor.commands.focus).not.toHaveBeenCalled();
+
+    flushFrame();
+    expect(editor.commands.focus).toHaveBeenCalledOnce();
+    // The retry loop re-checks on the next frame that focus stuck; since it
+    // did, it must not call focus() again.
+    flushFrame();
+    expect(editor.commands.focus).toHaveBeenCalledOnce();
+  });
+
+  it("retries on the next frame when a queued browser blur lands after the first restore", () => {
+    stubAnimationFrame();
+    const editor = makeEditor(true);
+    const { rerender } = renderHook(({ disabled }) => useSyncDisabledState(editor, disabled), {
+      initialProps: { disabled: false },
+    });
+
+    rerender({ disabled: true });
+    rerender({ disabled: false });
+
+    flushFrame();
+    expect(editor.commands.focus).toHaveBeenCalledOnce();
+
+    // A blur Chromium queued for the earlier disable lands only now,
+    // discarding the restore that already ran.
+    editor.simulateExternalBlur();
+
+    flushFrame();
+    expect(editor.commands.focus).toHaveBeenCalledTimes(2);
+  });
+
+  it("stops retrying once another control has since claimed focus", () => {
+    stubAnimationFrame();
+    const editor = makeEditor(true);
+    const claimant = document.createElement("input");
+    document.body.append(claimant);
+    const { rerender } = renderHook(({ disabled }) => useSyncDisabledState(editor, disabled), {
+      initialProps: { disabled: false },
+    });
+
+    rerender({ disabled: true });
+    claimant.focus();
+    rerender({ disabled: false });
+
+    flushFrame();
+    expect(editor.commands.focus).not.toHaveBeenCalled();
+    expect(frameQueue).toHaveLength(0);
+    claimant.remove();
+  });
+
+  it("does not focus on re-enable when the editor never had focus before disabling", () => {
+    stubAnimationFrame();
+    const editor = makeEditor(false);
+    const { rerender } = renderHook(({ disabled }) => useSyncDisabledState(editor, disabled), {
+      initialProps: { disabled: false },
+    });
+
+    rerender({ disabled: true });
+    rerender({ disabled: false });
+
+    expect(editor.commands.focus).not.toHaveBeenCalled();
+    expect(frameQueue).toHaveLength(0);
+  });
+
+  it("gives up after the retry budget instead of polling forever", () => {
+    stubAnimationFrame();
+    const editor = makeEditor(true);
+    // Focus never sticks and nothing else claims it either -- simulates a
+    // pathological case where the editor keeps losing focus every frame.
+    editor.commands.focus = vi.fn();
+    const { rerender } = renderHook(({ disabled }) => useSyncDisabledState(editor, disabled), {
+      initialProps: { disabled: false },
+    });
+
+    rerender({ disabled: true });
+    rerender({ disabled: false });
+
+    for (let i = 0; i < 10; i += 1) flushFrame();
+
+    expect(editor.commands.focus).toHaveBeenCalledTimes(5);
+    expect(frameQueue).toHaveLength(0);
+  });
+
+  it("stays armed on frame 1 when the queued browser blur has not landed yet, and restores focus once it does", () => {
+    stubAnimationFrame();
+    const editor = makeQueuedBlurEditor(true);
+    const { rerender } = renderHook(({ disabled }) => useSyncDisabledState(editor, disabled), {
+      initialProps: { disabled: false },
+    });
+
+    rerender({ disabled: true });
+    rerender({ disabled: false });
+
+    // Frame 1: the queued blur has not landed yet, so the editor still
+    // reports focus. That is not evidence the restore succeeded -- it is
+    // just that the browser hasn't gotten around to the blur it queued.
+    flushFrame();
+    expect(editor.commands.focus).not.toHaveBeenCalled();
+
+    // The queued blur lands only now.
+    editor.landQueuedBlur();
+
+    // The loop must still be armed to catch it.
+    flushFrame();
+    expect(editor.commands.focus).toHaveBeenCalledOnce();
+    expect(editor.view.hasFocus()).toBe(true);
+  });
+});
+
+describe("useSyncDisabledState cleanup", () => {
+  afterEach(() => {
+    vi.unstubAllGlobals();
+  });
+
+  it("stops retrying after unmount", () => {
+    stubAnimationFrame();
+    const editor = makeEditor(true);
+    editor.commands.focus = vi.fn();
+    const { rerender, unmount } = renderHook(
+      ({ disabled }) => useSyncDisabledState(editor, disabled),
+      { initialProps: { disabled: false } },
+    );
+
+    rerender({ disabled: true });
+    rerender({ disabled: false });
+
+    flushFrame();
+    expect(editor.commands.focus).toHaveBeenCalledOnce();
+
+    unmount();
+    flushFrame();
+    expect(editor.commands.focus).toHaveBeenCalledOnce();
+  });
+});
+
 describe("decideHistoryNav", () => {
   const base = {
     disabled: false,

apps/web/components/task/chat/use-tiptap-editor.ts

@@ -338,6 +338,46 @@ type SyncEditorOptions = {
   onChangeRef: React.RefObject<(value: string) => void>;
 };
 
+/** Editor methods required to synchronize the disabled state. */
+type DisabledStateEditor = {
+  view: { hasFocus: () => boolean };
+  setEditable: (editable: boolean) => void;
+  commands: { focus: () => void };
+};
+
+/** Bound focus restoration while browser focus settles after re-enabling. */
+const FOCUS_RESTORE_MAX_ATTEMPTS = 5;
+
+/** Preserve editor focus across temporary disabled states without stealing focus. */
+export function useSyncDisabledState(editor: DisabledStateEditor | null, disabled: boolean) {
+  const hadFocusBeforeDisableRef = useRef(false);
+  useEffect(() => {
+    if (!editor) return;
+    if (disabled) {
+      hadFocusBeforeDisableRef.current = editor.view.hasFocus();
+      editor.setEditable(false);
+      return;
+    }
+    editor.setEditable(true);
+    const hadFocus = hadFocusBeforeDisableRef.current;
+    hadFocusBeforeDisableRef.current = false;
+    if (!hadFocus) return;
+
+    let frame = 0;
+    let attempts = 0;
+    const tick = () => {
+      attempts += 1;
+      if (!editor.view.hasFocus()) {
+        if (!shouldRestoreFocusOnEnable(hadFocus)) return;
+        editor.commands.focus();
+      }
+      if (attempts < FOCUS_RESTORE_MAX_ATTEMPTS) frame = requestAnimationFrame(tick);
+    };
+    frame = requestAnimationFrame(tick);
+    return () => cancelAnimationFrame(frame);
+  }, [editor, disabled]);
+}
+
 function useSyncEditor({
   editor,
   disabled,
@@ -348,10 +388,7 @@ function useSyncEditor({
   initialSyncDoneRef,
   onChangeRef,
 }: SyncEditorOptions) {
-  // Sync disabled state
-  useEffect(() => {
-    if (editor) editor.setEditable(!disabled);
-  }, [editor, disabled]);
+  useSyncDisabledState(editor, disabled);
 
   // Sync placeholder via editor.storage. The DynamicPlaceholder extension reads
   // from editor.storage.dynamicPlaceholder.text at decoration time.
@@ -430,6 +467,15 @@ function syncEditorValue({
   initialSyncDoneRef.current = true;
 }
 
+// ── Focus-restore decision ───────────────────────────────────────────
+
+/** Return true only when the editor had focus and no other control owns focus. */
+export function shouldRestoreFocusOnEnable(hadFocusBeforeDisable: boolean): boolean {
+  if (!hadFocusBeforeDisable) return false;
+  const active = document.activeElement;
+  return active === null || active === document.body;
+}
+
 // ── Submit shortcut decision ────────────────────────────────────────
 
 export type SubmitShortcutDecision = "consume-noop" | "submit" | "defer";

apps/web/e2e/tests/chat/composer-focus-after-send-helpers.ts

@@ -0,0 +1,26 @@
+import { expect, type Page } from "@playwright/test";
+import type { SessionPage } from "../../pages/session-page";
+
+export async function assertComposerFocusAfterSend(
+  session: SessionPage,
+  page: Page,
+  submitFollowUp: () => Promise<void>,
+) {
+  const editor = session.activeChat().locator(".tiptap.ProseMirror:visible");
+
+  await session.sendMessageViaButton("first message after send");
+  await expect(
+    session.activeChat().getByText("first message after send", { exact: false }),
+  ).toBeVisible({ timeout: 15_000 });
+  await session.waitForChatIdle({ timeout: 30_000, requireEditable: true });
+  await expect(editor).toBeFocused({ timeout: 10_000 });
+
+  // No click or editor focus here. The composer must already hold focus.
+  await page.keyboard.type("second message after send");
+  await submitFollowUp();
+  await expect(
+    session.activeChat().getByText("second message after send", { exact: false }),
+  ).toBeVisible({ timeout: 15_000 });
+  await session.waitForChatIdle({ timeout: 30_000, requireEditable: true });
+  await expect(editor).toBeFocused({ timeout: 10_000 });
+}

apps/web/e2e/tests/chat/composer-focus-after-send.spec.ts

@@ -0,0 +1,32 @@
+import { test } from "../../fixtures/test-base";
+import { seedIdleSession } from "../../helpers/session";
+import { assertComposerFocusAfterSend } from "./composer-focus-after-send-helpers";
+
+// Regression coverage for the composer losing focus after a send. ProseMirror
+// maps its `editable` state onto the DOM `contenteditable` attribute, and a
+// real browser blurs the element when that attribute flips to `false` (which
+// every send does for its duration) without restoring focus when it flips
+// back -- jsdom cannot reproduce this, so Playwright is the only witness.
+test.describe("Composer focus after send", () => {
+  // Retries here absorb a pre-existing, unrelated dockview panel-portal race
+  // (the chat panel's React subtree occasionally remounts a moment after
+  // task load, independent of this composer's own send flow -- see the task
+  // plan for reproduction). A real regression in the focus-restore wiring
+  // fails on every retry, so this does not mask this spec's own assertion.
+  test.describe.configure({ retries: 1 });
+
+  test("keeps the composer focused across consecutive sends with no intervening click", async ({
+    testPage,
+    apiClient,
+    seedData,
+  }) => {
+    test.setTimeout(120_000);
+    const session = await seedIdleSession(
+      testPage,
+      apiClient,
+      seedData,
+      "Composer focus after send",
+    );
+    await assertComposerFocusAfterSend(session, testPage, () => session.clickSubmitWhenReady());
+  });
+});

apps/web/e2e/tests/chat/mobile-composer-focus-after-send.spec.ts

@@ -0,0 +1,24 @@
+import { test } from "../../fixtures/test-base";
+import { seedIdleSession } from "../../helpers/session";
+import { assertComposerFocusAfterSend } from "./composer-focus-after-send-helpers";
+
+// The mobile Playwright project uses touch input and only matches mobile-named
+// specs. Keep this regression on the real tap path, not only desktop click.
+test.describe("Mobile composer focus after send", () => {
+  test.describe.configure({ retries: 1 });
+
+  test("keeps the composer focused across consecutive touch sends", async ({
+    testPage,
+    apiClient,
+    seedData,
+  }) => {
+    test.setTimeout(120_000);
+    const session = await seedIdleSession(
+      testPage,
+      apiClient,
+      seedData,
+      "Mobile composer focus after send",
+    );
+    await assertComposerFocusAfterSend(session, testPage, () => session.tapSubmitWhenReady());
+  });
+});