From 63464c731cf5f01c14cbcf4e2b62623b737b3371 Mon Sep 17 00:00:00 2001
From: Willem Adnet <willem.adnet@grenoble-inp.org>
Date: Mon, 3 Aug 2026 17:26:25 +0200
Subject: [PATCH] feat(cli_ui): enable readline for the chat prompt and persist
 history

corridor_input() wraps its ANSI escapes in \x01/\x02, the readline markers for
non-printing characters, but nothing imported readline. Without it those two
bytes were written to the terminal verbatim, and the prompt had no history or
line editing.

Import readline (guarded, it is absent on some platforms) and load/save
~/.mlx_lm_chat_history around the chat session, on rank 0 only.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 mlx_lm/cli_ui.py   | 39 +++++++++++++++++++++++++++++++
 tests/test_chat.py | 57 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 96 insertions(+)

diff --git a/mlx_lm/cli_ui.py b/mlx_lm/cli_ui.py
index 133d1190e..1fe613fa3 100644
--- a/mlx_lm/cli_ui.py
+++ b/mlx_lm/cli_ui.py
@@ -1,11 +1,21 @@
 # Copyright © 2024 Apple Inc.
 
+import os
 import re
 import shutil
 import sys
 from contextlib import contextmanager
 from functools import lru_cache
 
+try:
+    # Importing readline is what makes input() honour the \x01/\x02
+    # non-printing markers corridor_input() emits, and what provides history
+    # and line editing. It is unavailable on some platforms (e.g. Windows
+    # without pyreadline), so the prompt must keep working without it.
+    import readline
+except ImportError:
+    readline = None
+
 import mlx.core as mx
 from rich.box import ROUNDED
 from rich.console import Console
@@ -249,6 +259,32 @@ def report_save(self, checkpoint):
         )
 
 
+CHAT_HISTORY_FILE = os.path.expanduser("~/.mlx_lm_chat_history")
+CHAT_HISTORY_LENGTH = 1000
+
+
+def load_chat_history(path: str = CHAT_HISTORY_FILE) -> None:
+    """Restore prompt history from a previous session, if there is one."""
+    if readline is None:
+        return
+    readline.set_history_length(CHAT_HISTORY_LENGTH)
+    try:
+        readline.read_history_file(path)
+    except (OSError, PermissionError):
+        # No history yet, or it is unreadable. Neither is worth failing over.
+        pass
+
+
+def save_chat_history(path: str = CHAT_HISTORY_FILE) -> None:
+    """Persist prompt history for the next session."""
+    if readline is None:
+        return
+    try:
+        readline.write_history_file(path)
+    except (OSError, PermissionError):
+        pass
+
+
 class ChatUI:
     """Helper class for rendering the chat UI and streaming responses."""
 
@@ -259,6 +295,7 @@ def __init__(self, args, rank: int = 0):
 
     def __enter__(self):
         if self._rank == 0:
+            load_chat_history()
             rows = [("model", str(self._args.model))]
             if self._args.adapter_path:
                 rows.append(("adapter", str(self._args.adapter_path)))
@@ -273,6 +310,8 @@ def __enter__(self):
         return self
 
     def __exit__(self, *exc):
+        if self._rank == 0:
+            save_chat_history()
         return False
 
     def prompt(self) -> str:
diff --git a/tests/test_chat.py b/tests/test_chat.py
index b4ddcf3ae..794ffff01 100644
--- a/tests/test_chat.py
+++ b/tests/test_chat.py
@@ -1,7 +1,10 @@
 import argparse
+import os
+import tempfile
 import unittest
 from unittest.mock import MagicMock, patch
 
+from mlx_lm import cli_ui
 from mlx_lm.chat import setup_arg_parser
 
 
@@ -171,5 +174,59 @@ def test_no_system_prompt_in_messages(
         self.assertEqual(call_args[0]["content"], "What is the weather?")
 
 
+class TestChatHistory(unittest.TestCase):
+    """Prompt history is persisted across sessions, on rank 0 only."""
+
+    def setUp(self):
+        self.tmp = tempfile.TemporaryDirectory()
+        self.path = os.path.join(self.tmp.name, "history")
+
+    def tearDown(self):
+        self.tmp.cleanup()
+
+    def test_round_trip(self):
+        if cli_ui.readline is None:
+            self.skipTest("readline unavailable on this platform")
+        cli_ui.readline.clear_history()
+        cli_ui.readline.add_history("hello there")
+        cli_ui.save_chat_history(self.path)
+
+        cli_ui.readline.clear_history()
+        self.assertEqual(cli_ui.readline.get_current_history_length(), 0)
+
+        cli_ui.load_chat_history(self.path)
+        self.assertEqual(cli_ui.readline.get_history_item(1), "hello there")
+
+    def test_missing_file_is_not_an_error(self):
+        cli_ui.load_chat_history(os.path.join(self.tmp.name, "does-not-exist"))
+
+    def test_unwritable_path_is_not_an_error(self):
+        cli_ui.save_chat_history(os.path.join(self.tmp.name, "no", "such", "dir"))
+
+    def test_history_is_skipped_without_readline(self):
+        with patch.object(cli_ui, "readline", None):
+            cli_ui.load_chat_history(self.path)
+            cli_ui.save_chat_history(self.path)
+        self.assertFalse(os.path.exists(self.path))
+
+    def test_only_rank_zero_touches_history(self):
+        args = MagicMock()
+        args.model, args.adapter_path = "m", None
+        args.max_tokens, args.system_prompt = 10, None
+
+        with patch.object(cli_ui, "load_chat_history") as load, patch.object(
+            cli_ui, "save_chat_history"
+        ) as save:
+            with cli_ui.ChatUI(args, rank=1):
+                pass
+            load.assert_not_called()
+            save.assert_not_called()
+
+            with cli_ui.ChatUI(args, rank=0):
+                pass
+            load.assert_called_once()
+            save.assert_called_once()
+
+
 if __name__ == "__main__":
     unittest.main()