Enhancement: Interactive Shell with Syntax Highlighting

src/documentdb/shell/DocumentDBShellPty.ts

@@ -9,6 +9,7 @@ import { ext } from '../../extensionVariables';
 import { deserializeResultForSchema, feedResultToSchemaStore } from '../feedResultToSchemaStore';
 import { type SerializableExecutionResult } from '../playground/workerTypes';
 import { SchemaStore } from '../SchemaStore';
+import { colorizeShellInput } from './highlighting/colorizeShellInput';
 import { SettingsHintError } from './SettingsHintError';
 import { type CompletionResult, ShellCompletionProvider } from './ShellCompletionProvider';
 import { findCommonPrefix, renderCompletionList } from './ShellCompletionRenderer';
@@ -132,6 +133,12 @@ export class DocumentDBShellPty implements vscode.Pseudoterminal {
             onLine: (line: string) => void this.handleLineInput(line),
             onInterrupt: () => this.handleInterrupt(),
             onContinuation: () => this.showContinuationPrompt(),
+            colorize: (input: string) => {
+                if (!this.isColorEnabled()) {
+                    return input;
+                }
+                return colorizeShellInput(input);
+            },
             onTab: (buffer: string, cursor: number) => this.handleTab(buffer, cursor),
             onBufferChange: (buffer: string, cursor: number) => this.handleBufferChange(buffer, cursor),
             onAcceptGhostText: () => this.handleAcceptGhostText(),
@@ -141,9 +148,10 @@ export class DocumentDBShellPty implements vscode.Pseudoterminal {
     // ─── Pseudoterminal interface ────────────────────────────────────────────
 
     open(initialDimensions: vscode.TerminalDimensions | undefined): void {
-        // Track terminal width for completion rendering
+        // Track terminal width for completion rendering and wrap-aware re-rendering
         if (initialDimensions) {
             this._columns = initialDimensions.columns;
+            this._inputHandler.setColumns(initialDimensions.columns);
         }
 
         // Disable input during initialization to prevent race conditions
@@ -222,6 +230,7 @@ export class DocumentDBShellPty implements vscode.Pseudoterminal {
      */
     setDimensions(dimensions: vscode.TerminalDimensions): void {
         this._columns = dimensions.columns;
+        this._inputHandler.setColumns(dimensions.columns);
     }
 
     // ─── Private: Session initialization ─────────────────────────────────────
@@ -701,18 +710,13 @@ export class DocumentDBShellPty implements vscode.Pseudoterminal {
 
     /**
      * Rewrite the prompt and current buffer after showing a completion list.
+     * Uses the colorize callback (via renderCurrentLine) so highlighting is preserved.
      */
     private rewriteCurrentLine(): void {
         const prompt = `${this._currentDatabase}> `;
-        const buffer = this._inputHandler.getBuffer();
-        const cursor = this._inputHandler.getCursor();
-        this._writeEmitter.fire(prompt + buffer);
-
-        // Position cursor at the correct location
-        const trailingChars = buffer.length - cursor;
-        if (trailingChars > 0) {
-            this._writeEmitter.fire(`\x1b[${String(trailingChars)}D`);
-        }
+        this._inputHandler.setPromptWidth(prompt.length);
+        this._writeEmitter.fire(prompt);
+        this._inputHandler.renderCurrentLine();
     }
 
     // ─── Private: Ghost text ─────────────────────────────────────────────────

src/documentdb/shell/ShellGhostText.ts

@@ -15,6 +15,8 @@
  * continues typing at the same location.
  */
 
+import { terminalDisplayWidth } from './terminalDisplayWidth';
+
 // ─── ANSI constants ──────────────────────────────────────────────────────────
 
 /** Dim + gray for ghost text appearance. */
@@ -23,62 +25,6 @@ const ANSI_RESET = '\x1b[0m';
 /** Erase from cursor to end of line. */
 const ERASE_TO_EOL = '\x1b[K';
 
-/**
- * Compute the display-column width of a string for the terminal.
- *
- * JavaScript's `String.length` counts UTF-16 code units, but ANSI cursor
- * movement operates on display columns. Surrogate pairs (emoji, symbols
- * above U+FFFF) are 2 code units but typically 1–2 terminal columns.
- *
- * This uses `Intl.Segmenter` (available in Node 16+) to iterate grapheme
- * clusters and counts each one as 1 column unless it is a known
- * full-width/wide character (CJK Unified Ideographs, etc.).
- */
-function terminalDisplayWidth(text: string): number {
-    // Fast path: ASCII-only strings (common case)
-    if (/^[\x20-\x7e]*$/.test(text)) {
-        return text.length;
-    }
-
-    let width = 0;
-    // Use Intl.Segmenter to properly iterate grapheme clusters
-    // This handles surrogate pairs, combining marks, ZWJ sequences, etc.
-    const segmenter = new Intl.Segmenter(undefined, { granularity: 'grapheme' });
-    for (const { segment } of segmenter.segment(text)) {
-        const cp = segment.codePointAt(0) ?? 0;
-        // Full-width / wide characters occupy 2 columns
-        if (isWideCharacter(cp)) {
-            width += 2;
-        } else {
-            width += 1;
-        }
-    }
-
-    return width;
-}
-
-/**
- * Returns true for code points that occupy 2 terminal columns.
- * Covers CJK Unified Ideographs and common full-width ranges.
- */
-function isWideCharacter(cp: number): boolean {
-    return (
-        (cp >= 0x1100 && cp <= 0x115f) || // Hangul Jamo
-        (cp >= 0x2e80 && cp <= 0x303e) || // CJK Radicals, Kangxi, CJK Symbols
-        (cp >= 0x3040 && cp <= 0x33bf) || // Hiragana, Katakana, Bopomofo, etc.
-        (cp >= 0x3400 && cp <= 0x4dbf) || // CJK Unified Ideographs Extension A
-        (cp >= 0x4e00 && cp <= 0xa4cf) || // CJK Unified Ideographs + Yi
-        (cp >= 0xa960 && cp <= 0xa97c) || // Hangul Jamo Extended-A
-        (cp >= 0xac00 && cp <= 0xd7a3) || // Hangul Syllables
-        (cp >= 0xf900 && cp <= 0xfaff) || // CJK Compatibility Ideographs
-        (cp >= 0xfe30 && cp <= 0xfe6f) || // CJK Compatibility Forms
-        (cp >= 0xff01 && cp <= 0xff60) || // Fullwidth Forms
-        (cp >= 0xffe0 && cp <= 0xffe6) || // Fullwidth Signs
-        (cp >= 0x20000 && cp <= 0x2fffd) || // CJK Unified Ideographs Extension B+
-        (cp >= 0x30000 && cp <= 0x3fffd) // CJK Unified Ideographs Extension G+
-    );
-}
-
 /**
  * Manages the lifecycle of ghost text in the terminal.
  *

src/documentdb/shell/ShellInputHandler.test.ts

@@ -40,8 +40,9 @@ describe('ShellInputHandler', () => {
     describe('basic character input', () => {
         it('should echo printable characters', () => {
             handler.handleInput('a');
-            expect(written).toBe('a');
             expect(handler.getBuffer()).toBe('a');
+            // The re-render output contains ANSI positioning + the character
+            expect(written).toContain('a');
         });
 
         it('should accumulate characters in buffer', () => {
@@ -56,7 +57,8 @@ describe('ShellInputHandler', () => {
         it('should handle multi-character input at once', () => {
             handler.handleInput('hello');
             expect(handler.getBuffer()).toBe('hello');
-            expect(written).toBe('hello');
+            // The re-render output contains ANSI positioning + the full text
+            expect(written).toContain('hello');
         });
 
         it('should ignore control characters below space (except special ones)', () => {