fix(prompt-input): prevent extra line break after send (#136)

Author: heath-s

src/components/ai-elements/prompt-lexical-editor.commands.ts

@@ -0,0 +1,20 @@
+import {
+  COMMAND_PRIORITY_CRITICAL,
+  KEY_ENTER_COMMAND,
+  type LexicalEditor,
+} from "lexical";
+
+/**
+ * The prompt composer handles plain Enter during React's capture phase.
+ * Lexical's plain-text plugin does not honor `defaultPrevented` and would
+ * otherwise insert a trailing line break after the prompt has been sent.
+ */
+export function registerPromptLexicalPreventedEnterCommand(
+  editor: LexicalEditor,
+) {
+  return editor.registerCommand(
+    KEY_ENTER_COMMAND,
+    (event) => event?.defaultPrevented === true,
+    COMMAND_PRIORITY_CRITICAL,
+  );
+}

src/components/ai-elements/prompt-lexical-editor.tsx

@@ -54,6 +54,9 @@ import {
 } from "@/lib/prompt-token-chips";
 import type { WorkspaceInformationReferenceOption } from "@/lib/workspace-information-references";
 import { cn } from "@/lib/utils";
+import {
+  registerPromptLexicalPreventedEnterCommand,
+} from "./prompt-lexical-editor.commands";
 import { PromptTokenChip } from "./prompt-token-chip";
 
 const PROMPT_SYNC_TAG = "stave-prompt-sync";
@@ -576,6 +579,17 @@ function PromptLexicalTokenDeletionPlugin() {
   return null;
 }
 
+function PromptLexicalPreventedEnterPlugin() {
+  const [editor] = useLexicalComposerContext();
+
+  useEffect(
+    () => registerPromptLexicalPreventedEnterCommand(editor),
+    [editor],
+  );
+
+  return null;
+}
+
 function PromptLexicalImperativePlugin(args: {
   forwardedRef: ForwardedRef<PromptLexicalEditorHandle>;
 }) {
@@ -783,6 +797,7 @@ export const PromptLexicalEditor = forwardRef<
       <HistoryPlugin />
       <PromptLexicalEditablePlugin disabled={props.disabled} />
       <PromptLexicalSelectionPlugin onSelectionChange={props.onSelectionChange} />
+      <PromptLexicalPreventedEnterPlugin />
       <PromptLexicalTokenDeletionPlugin />
       <PromptLexicalImperativePlugin forwardedRef={ref} />
       <PromptLexicalExternalSyncPlugin

tests/prompt-lexical-editor.commands.test.ts

@@ -0,0 +1,59 @@
+import { describe, expect, test } from "bun:test";
+import {
+  COMMAND_PRIORITY_EDITOR,
+  createEditor,
+  KEY_ENTER_COMMAND,
+} from "lexical";
+import {
+  registerPromptLexicalPreventedEnterCommand,
+} from "@/components/ai-elements/prompt-lexical-editor.commands";
+
+describe("prompt Lexical Enter handling", () => {
+  test("consumes Enter after the prompt composer prevents it", () => {
+    const editor = createEditor({ namespace: "prompt-enter-prevented" });
+    let plainTextHandlerCalls = 0;
+    const unregisterGuard =
+      registerPromptLexicalPreventedEnterCommand(editor);
+    const unregisterPlainTextHandler = editor.registerCommand(
+      KEY_ENTER_COMMAND,
+      () => {
+        plainTextHandlerCalls += 1;
+        return true;
+      },
+      COMMAND_PRIORITY_EDITOR,
+    );
+
+    const handled = editor.dispatchCommand(KEY_ENTER_COMMAND, {
+      defaultPrevented: true,
+    } as KeyboardEvent);
+
+    expect(handled).toBe(true);
+    expect(plainTextHandlerCalls).toBe(0);
+    unregisterPlainTextHandler();
+    unregisterGuard();
+  });
+
+  test("leaves multiline Enter behavior to Lexical when it is not prevented", () => {
+    const editor = createEditor({ namespace: "prompt-enter-multiline" });
+    let plainTextHandlerCalls = 0;
+    const unregisterGuard =
+      registerPromptLexicalPreventedEnterCommand(editor);
+    const unregisterPlainTextHandler = editor.registerCommand(
+      KEY_ENTER_COMMAND,
+      () => {
+        plainTextHandlerCalls += 1;
+        return true;
+      },
+      COMMAND_PRIORITY_EDITOR,
+    );
+
+    const handled = editor.dispatchCommand(KEY_ENTER_COMMAND, {
+      defaultPrevented: false,
+    } as KeyboardEvent);
+
+    expect(handled).toBe(true);
+    expect(plainTextHandlerCalls).toBe(1);
+    unregisterPlainTextHandler();
+    unregisterGuard();
+  });
+});