kk

CLI.py

@@ -30,12 +30,16 @@ try:
     from prompt_toolkit import PromptSession
     from prompt_toolkit.completion import Completer, Completion
     from prompt_toolkit.document import Document
+    from prompt_toolkit.lexers import Lexer
+    from prompt_toolkit.styles import Style
     PROMPT_TOOLKIT_AVAILABLE = True
 except ImportError:  # pragma: no cover - optional dependency
     PromptSession = None  # type: ignore
     Completer = None  # type: ignore
     Completion = None  # type: ignore
     Document = None  # type: ignore
+    Lexer = None  # type: ignore
+    Style = None  # type: ignore
     PROMPT_TOOLKIT_AVAILABLE = False
 
 
@@ -531,6 +535,46 @@ if (
         async def get_completions_async(self, document: Document, complete_event):  # type: ignore[override]
             for completion in self.get_completions(document, complete_event):
                 yield completion
+
+    class MedeiaLexer(Lexer):
+        def lex_document(self, document):
+            def get_line(lineno):
+                line = document.lines[lineno]
+                tokens = []
+                
+                import re
+                # Match: Whitespace, Pipe, Quoted string, or Word
+                pattern = re.compile(r'''
+                    (\s+) |                                      # 1. Whitespace
+                    (\|) |                                       # 2. Pipe
+                    ("(?:[^"\\]|\\.)*"|'(?:[^'\\]|\\.)*') |      # 3. Quoted string
+                    ([^\s\|]+)                                   # 4. Word
+                ''', re.VERBOSE)
+                
+                is_cmdlet = True
+                
+                for match in pattern.finditer(line):
+                    ws, pipe, quote, word = match.groups()
+                    
+                    if ws:
+                        tokens.append(('', ws))
+                    elif pipe:
+                        tokens.append(('class:pipe', pipe))
+                        is_cmdlet = True
+                    elif quote:
+                        tokens.append(('class:string', quote))
+                        is_cmdlet = False
+                    elif word:
+                        if is_cmdlet:
+                            tokens.append(('class:cmdlet', word))
+                            is_cmdlet = False
+                        elif word.startswith('-'):
+                            tokens.append(('class:argument', word))
+                        else:
+                            tokens.append(('class:value', word))
+                            
+                return tokens
+            return get_line
 else:  # pragma: no cover - prompt toolkit unavailable
     CmdletCompleter = None  # type: ignore[assignment]
 
@@ -586,7 +630,21 @@ Example:  search-file --help
         
         if PROMPT_TOOLKIT_AVAILABLE and PromptSession is not None and CmdletCompleter is not None:
             completer = CmdletCompleter()
-            session = PromptSession(completer=cast(Any, completer))
+            
+            # Define style for syntax highlighting
+            style = Style.from_dict({
+                'cmdlet': '#ffffff',   # white
+                'argument': '#3b8eea', # blue-ish
+                'value': '#ce9178',    # red-ish
+                'string': '#ce55ff',   # purple
+                'pipe': '#4caf50',     # green
+            })
+            
+            session = PromptSession(
+                completer=cast(Any, completer),
+                lexer=MedeiaLexer(),
+                style=style
+            )
 
             def get_input(prompt: str = ">>>|") -> str:
                 return session.prompt(prompt)
@@ -645,6 +703,7 @@ Example:  search-file --help
                         if last_table is None:
                             last_table = ctx.get_last_result_table()
                             
+
                         if last_table:
                             print()
                             # Also update current stage table so @N expansion works correctly
@@ -779,10 +838,10 @@ def _execute_pipeline(tokens: list):
             else:
                 # Try command-based expansion first if we have source command info
                 command_expanded = False
+                selected_row_args = []
                 
                 if source_cmd:
                     # Try to find row args for the selected indices
-                    selected_row_args = []
                     for idx in first_stage_selection_indices:
                         row_args = ctx.get_current_stage_table_row_selection_args(idx)
                         if row_args: