Address PR feedback and fix minor issues with selection

Author: ernst-dev

pages/prompt-input/token-renderer.page.tsx

@@ -5,33 +5,11 @@ import ReactDOM from 'react-dom';
 
 import { PromptInputProps } from '~components/prompt-input';
 import { extractTokensFromDOM } from '~components/prompt-input/core/token-operations';
-import { PortalContainer, RenderTokenProps, renderTokensToDOM } from '~components/prompt-input/core/token-renderer';
+import { PortalContainer, renderTokensToDOM } from '~components/prompt-input/core/token-renderer';
+import Token from '~components/token/internal';
 
 import { SimplePage } from '../app/templates';
 
-// Custom token renderer — intentionally NOT using the Token component.
-// This proves the renderer is decoupled from any specific UI component.
-function CustomToken({ label }: RenderTokenProps) {
-  return (
-    <span
-      style={{
-        display: 'inline-block',
-        padding: '2px 8px',
-        margin: '0 2px',
-        borderRadius: 12,
-        background: '#fce4ec',
-        border: '2px dashed #e91e63',
-        fontSize: 13,
-        lineHeight: '18px',
-        fontWeight: 600,
-        color: '#880e4f',
-      }}
-    >
-      ⚡ {label}
-    </span>
-  );
-}
-
 let nextId = 1;
 
 // Menu definitions for trigger/reference token extraction
@@ -158,16 +136,9 @@ export default function TokenRendererPage() {
         }}
       />
 
-      {/* Render reference tokens into their DOM containers via portals */}
       {Array.from(portalContainersRef.current.values()).map(container =>
         ReactDOM.createPortal(
-          <CustomToken
-            key={container.id}
-            id={container.id}
-            label={container.label}
-            disabled={false}
-            readOnly={false}
-          />,
+          <Token key={container.id} variant="inline" label={container.label} disabled={false} readOnly={false} />,
           container.element
         )
       )}

src/prompt-input/__integ__/prompt-input-token-mode.test.ts

@@ -282,18 +282,18 @@ const setupTest = (testFn: (page: PromptInputTokenModePage) => Promise<void>) =>
       await page.keys(['Home']);
       await page.pause(100);
 
-      // Select forward: "hi " then the reference then " bye"
+      // Select "hi " (3 characters)
       for (let i = 0; i < 3; i++) {
         await page.keys(['Shift', 'ArrowRight', 'Shift']);
       }
       expect(await page.getSelectedText()).toBe('hi ');
 
-      // One more should jump over the reference
+      // One more jumps over the atomic reference
       await page.keys(['Shift', 'ArrowRight', 'Shift']);
       await page.pause(100);
       const selected = await page.getSelectedText();
       expect(selected).toContain('hi ');
-      expect(selected).toContain('John Smith');
+      expect(selected).toContain('Jane Smith');
     })
   );
 
@@ -309,24 +309,24 @@ const setupTest = (testFn: (page: PromptInputTokenModePage) => Promise<void>) =>
       await page.keys([' ', 'b', 'y', 'e']);
       await page.pause(200);
 
-      // Cursor is at end. Select backward: "bye " then reference then " hi"
+      // Select backward from end: " bye" (4 chars)
       for (let i = 0; i < 4; i++) {
         await page.keys(['Shift', 'ArrowLeft', 'Shift']);
       }
       const afterText = await page.getSelectedText();
       expect(afterText).toBe(' bye');
 
-      // One more should jump over the reference
+      // One more jumps over the atomic reference
       await page.keys(['Shift', 'ArrowLeft', 'Shift']);
       await page.pause(100);
       const selected = await page.getSelectedText();
-      expect(selected).toContain('John Smith');
+      expect(selected).toContain('Jane Smith');
       expect(selected).toContain(' bye');
     })
   );
 
   test(
-    'shift+left then shift+right reversal deselects reference without flipping selection',
+    'shift+left then shift+right reversal deselects correctly around reference',
     setupTest(async page => {
       await page.focusInput();
       await page.keys(['h', 'e', 'l', 'l', 'o', ' ']);
@@ -337,27 +337,121 @@ const setupTest = (testFn: (page: PromptInputTokenModePage) => Promise<void>) =>
       await page.keys([' ', 'w', 'o', 'r', 'l', 'd']);
       await page.pause(200);
 
-      // Place cursor in middle of "world" (3 chars from end)
+      // Place cursor in middle of "world"
       await page.keys(['ArrowLeft', 'ArrowLeft', 'ArrowLeft']);
       await page.pause(100);
 
-      // Select backward past " wo", over reference, into "hello"
-      for (let i = 0; i < 7; i++) {
+      // Select backward: " wo" (3) + reference (1) + "hello " (6) = 10 presses
+      for (let i = 0; i < 10; i++) {
         await page.keys(['Shift', 'ArrowLeft', 'Shift']);
       }
       await page.pause(100);
       const backwardSel = await page.getSelectedText();
-      expect(backwardSel).toContain('John Smith');
+      expect(backwardSel).toContain('hello');
+      expect(backwardSel).toContain('Jane Smith');
 
-      // Now reverse with shift+right — deselect back through "hello " and the reference
-      for (let i = 0; i < 7; i++) {
+      // Reverse with shift+right — deselect everything
+      for (let i = 0; i < 10; i++) {
         await page.keys(['Shift', 'ArrowRight', 'Shift']);
       }
       await page.pause(100);
 
-      // Selection should be collapsed or very small — not extending the wrong end
       const afterReverse = await page.getSelectedText();
-      expect(afterReverse.length).toBeLessThanOrEqual(1);
+      expect(afterReverse).toBe('');
+    })
+  );
+});
+
+(isReact18 ? describe : describe.skip)('PromptInput token mode - delete key with references', () => {
+  test(
+    'delete key removes reference token ahead of cursor',
+    setupTest(async page => {
+      await page.focusInput();
+      await page.keys(['@']);
+      await page.pause(200);
+      await page.keys(['ArrowDown', 'Enter']);
+      await page.pause(200);
+      await page.keys([' ', 'h', 'i']);
+      await page.pause(100);
+
+      // Move cursor to start (before the reference)
+      await page.keys(['Home']);
+      await page.pause(100);
+
+      // Delete should remove the reference
+      await page.keys(['Delete']);
+      await page.pause(200);
+
+      const text = await page.getEditorText();
+      expect(text.trim()).toBe('hi');
+      expect(text).not.toContain('Jane Smith');
+    })
+  );
+
+  test(
+    'backspace removes reference token behind cursor',
+    setupTest(async page => {
+      await page.focusInput();
+      await page.keys(['h', 'i', ' ']);
+      await page.keys(['@']);
+      await page.pause(200);
+      await page.keys(['ArrowDown', 'Enter']);
+      await page.pause(200);
+
+      // Cursor is right after the reference — backspace removes it
+      await page.keys(['Backspace']);
+      await page.pause(200);
+
+      const text = await page.getEditorText();
+      expect(text.trim()).toBe('hi');
+      expect(text).not.toContain('Jane Smith');
+      expect(await page.getCaretOffset()).toBe(3);
+    })
+  );
+});
+
+(isReact18 ? describe : describe.skip)('PromptInput token mode - trigger dismissal', () => {
+  test(
+    'space on empty trigger dismisses it',
+    setupTest(async page => {
+      await page.focusInput();
+      await page.keys(['h', 'i', ' ', '@']);
+      await page.pause(200);
+      await expect(page.isMenuOpen()).resolves.toBe(true);
+
+      await page.keys([' ']);
+      await page.pause(300);
+
+      await expect(page.isMenuOpen()).resolves.toBe(false);
+      const text = await page.getEditorText();
+      expect(text).toContain('hi');
+      expect(text).toContain('@');
+    })
+  );
+});
+
+(isReact18 ? describe : describe.skip)('PromptInput token mode - multiple references', () => {
+  test(
+    'insert two references with text between them',
+    setupTest(async page => {
+      await page.focusInput();
+      await page.keys(['@']);
+      await page.pause(200);
+      await page.keys(['ArrowDown', 'Enter']);
+      await page.pause(200);
+
+      await page.keys([' ', 'a', 'n', 'd', ' ']);
+
+      await page.keys(['@']);
+      await page.pause(200);
+      // Select second option
+      await page.keys(['ArrowDown', 'ArrowDown', 'Enter']);
+      await page.pause(200);
+
+      const text = await page.getEditorText();
+      expect(text).toContain('Jane Smith');
+      expect(text).toContain('and');
+      expect(text).toContain('Bob');
     })
   );
 });

src/prompt-input/__integ__/token-renderer.test.ts

@@ -63,7 +63,6 @@ describe('Token Renderer (isolated)', () => {
       await clickButton(page, 'Add reference');
 
       const editorText = await getEditorText(page);
-      expect(editorText).toContain('⚡');
       expect(editorText).toContain('Alice');
 
       const tokenState = await getTokenState(page);

src/prompt-input/__tests__/caret-controller.test.ts

@@ -3,6 +3,7 @@
 
 jest.mock('../styles.css.js', () => ({}), { virtual: true });
 
+import './jsdom-polyfills';
 import {
   calculateTokenPosition,
   calculateTotalTokenLength,
@@ -783,11 +784,11 @@ describe('CaretController - additional branch coverage', () => {
       expect(window.getSelection()!.toString()).toBe('');
     });
 
-    test('does nothing when no selection object', () => {
+    test('selectAll selects content even when element is not focused', () => {
       addParagraph(el, 'hello');
       controller.selectAll();
-      // selectAll should work without throwing
-      expect(window.getSelection()!.rangeCount).toBeGreaterThanOrEqual(0);
+      const sel = window.getSelection()!;
+      expect(sel.toString()).toBe('hello');
     });
   });
 
@@ -912,8 +913,8 @@ describe('CaretController - additional branch coverage', () => {
       window.getSelection()?.addRange(range);
 
       const pos = controller.getPosition();
-      // Should be reference (1) + offset in after content
-      expect(pos).toBeGreaterThanOrEqual(1);
+      // Reference (1) + stripped offset in after-spot: raw offset 3 minus 1 for zero-width char = 2
+      expect(pos).toBe(3);
     });
 
     test('getPosition with cursor in trigger text node', () => {
@@ -2073,7 +2074,6 @@ describe('CaretController - defensive guard coverage', () => {
   afterEach(() => {
     document.body.innerHTML = '';
     jest.restoreAllMocks();
-    delete (Range.prototype as any).getBoundingClientRect;
     delete (HTMLElement.prototype as any).scrollIntoView;
   });
 
@@ -2104,7 +2104,7 @@ describe('CaretController - defensive guard coverage', () => {
     const mockRangeRect = { top: 150, bottom: 160, left: 0, right: 10 };
 
     jest.spyOn(el, 'getBoundingClientRect').mockReturnValue(mockElementRect as DOMRect);
-    Range.prototype.getBoundingClientRect = jest.fn().mockReturnValue(mockRangeRect as DOMRect);
+    jest.spyOn(Range.prototype, 'getBoundingClientRect').mockReturnValue(mockRangeRect as DOMRect);
 
     const scrollSpy = jest.fn();
     HTMLElement.prototype.scrollIntoView = scrollSpy;
@@ -2121,27 +2121,14 @@ describe('CaretController - defensive guard coverage', () => {
     const mockRangeRect = { top: 150, bottom: 160, left: 0, right: 10 };
 
     jest.spyOn(el, 'getBoundingClientRect').mockReturnValue(mockElementRect as DOMRect);
-    Range.prototype.getBoundingClientRect = jest.fn().mockReturnValue(mockRangeRect as DOMRect);
+    jest.spyOn(Range.prototype, 'getBoundingClientRect').mockReturnValue(mockRangeRect as DOMRect);
     const scrollSpy = jest.fn();
     HTMLElement.prototype.scrollIntoView = scrollSpy;
 
     controller.setPosition(2, 8);
     expect(scrollSpy).toHaveBeenCalled();
   });
 
-  test('setPosition scroll handles getBoundingClientRect throwing', () => {
-    addParagraph(el, 'hello');
-    el.focus();
-
-    Range.prototype.getBoundingClientRect = jest.fn().mockImplementation(() => {
-      throw new Error('not supported');
-    });
-
-    // Should not throw — caught by try/catch, and position should still be set
-    controller.setPosition(3);
-    expect(controller.getPosition()).toBe(3);
-  });
-
   test('restore does nothing when element is not the active element', () => {
     addParagraph(el, 'hello');
     el.focus();
@@ -2480,7 +2467,6 @@ describe('CaretController - remaining uncovered branches', () => {
   afterEach(() => {
     document.body.innerHTML = '';
     jest.restoreAllMocks();
-    delete (Range.prototype as any).getBoundingClientRect;
     delete (HTMLElement.prototype as any).scrollIntoView;
   });
 
@@ -2492,7 +2478,7 @@ describe('CaretController - remaining uncovered branches', () => {
     const mockRangeRect = { top: 150, bottom: 160, left: 0, right: 10 };
 
     jest.spyOn(el, 'getBoundingClientRect').mockReturnValue(mockElementRect as DOMRect);
-    Range.prototype.getBoundingClientRect = jest.fn().mockReturnValue(mockRangeRect as DOMRect);
+    jest.spyOn(Range.prototype, 'getBoundingClientRect').mockReturnValue(mockRangeRect as DOMRect);
     HTMLElement.prototype.scrollIntoView = jest.fn();
 
     // Start at 0 (valid), end at 999 (beyond content)
@@ -2509,7 +2495,7 @@ describe('CaretController - remaining uncovered branches', () => {
     const mockRangeRect = { top: 150, bottom: 160, left: 0, right: 10 };
 
     jest.spyOn(el, 'getBoundingClientRect').mockReturnValue(mockElementRect as DOMRect);
-    Range.prototype.getBoundingClientRect = jest.fn().mockReturnValue(mockRangeRect as DOMRect);
+    jest.spyOn(Range.prototype, 'getBoundingClientRect').mockReturnValue(mockRangeRect as DOMRect);
     HTMLElement.prototype.scrollIntoView = jest.fn();
 
     // Collapsed range (no end)
@@ -2699,10 +2685,12 @@ describe('CaretController - null textContent fallbacks', () => {
 
     Object.defineProperty(textNode, 'textContent', { value: null, writable: true });
 
+    el.focus();
     controller.positionAfterText(textNode);
-    // Should not throw — falls back to offset 0
+    // Falls back to offset 0 since textContent is null
     const sel = window.getSelection()!;
-    expect(sel.rangeCount).toBeGreaterThanOrEqual(0);
+    expect(sel.rangeCount).toBe(1);
+    expect(sel.getRangeAt(0).startOffset).toBe(0);
   });
 });