feat: Add option to save commands to bash history

llm_cmd.py

@@ -1,5 +1,6 @@
 import click
 import llm
+import os
 import subprocess
 from prompt_toolkit import PromptSession
 from prompt_toolkit.lexers import PygmentsLexer
@@ -21,7 +22,8 @@ def register_commands(cli):
     @click.option("-m", "--model", default=None, help="Specify the model to use")
     @click.option("-s", "--system", help="Custom system prompt")
     @click.option("--key", help="API key to use")
-    def cmd(args, model, system, key):
+    @click.option("--save-history", is_flag=True, help="Save commands to shell history")
+    def cmd(args, model, system, key, save_history):
         """Generate and execute commands in your shell"""
         from llm.cli import get_default_model
         prompt = " ".join(args)
@@ -30,9 +32,11 @@ def cmd(args, model, system, key):
         if model_obj.needs_key:
             model_obj.key = llm.get_key(key, model_obj.needs_key, model_obj.key_env_var)
         result = model_obj.prompt(prompt, system=system or SYSTEM_PROMPT)
-        interactive_exec(str(result))
+        interactive_exec(str(result), save_history)
 
-def interactive_exec(command):
+def interactive_exec(command, save_history):
+    # Check if history saving is enabled via flag or env var
+    save_to_history = save_history or os.environ.get('LLM_CMD_SAVE_HISTORY')
     session = PromptSession(lexer=PygmentsLexer(BashLexer))
     with patch_stdout():
         if '\n' in command:
@@ -41,9 +45,29 @@ def interactive_exec(command):
         else:
             edited_command = session.prompt("> ", default=command)
     try:
+        # Execute the command first
         output = subprocess.check_output(
             edited_command, shell=True, stderr=subprocess.STDOUT
         )
         print(output.decode())
+
+        # Only save successful commands to history
+        if save_to_history:
+            histfile = os.environ.get('HISTFILE')
+            if not histfile:
+                print("Warning: $HISTFILE environment variable not set or not exported")
+                print("History saving is disabled. Please ensure HISTFILE is exported in your shell config")
+            else:
+                # Add to bash history and append to history file
+                escaped_cmd = edited_command.replace('"', '\\"')
+                history_result = subprocess.run(
+                    ['bash', '-c', f'history -s "{escaped_cmd}" && history -a'],
+                    capture_output=True,
+                    text=True,
+                    check=False
+                )
+                if history_result.returncode != 0:
+                    print("Warning: Failed to save command to shell history")
     except subprocess.CalledProcessError as e:
-        print(f"Command failed with error (exit status {e.returncode}): {e.output.decode()}")
+        error_msg = e.output.decode() if e.output else "No error output available"
+        print(f"Command failed with error (exit status {e.returncode}): {error_msg}")

tests/test_history_saving.py

@@ -0,0 +1,83 @@
+import os
+import pytest
+from unittest.mock import patch, MagicMock
+from llm_cmd import interactive_exec
+
+@pytest.fixture
+def mock_subprocess_run():
+    with patch('subprocess.run') as mock_run:
+        yield mock_run
+
+@pytest.fixture
+def mock_subprocess_check_output():
+    with patch('subprocess.check_output') as mock_check:
+        yield mock_check
+
+@pytest.fixture
+def mock_prompt_session():
+    with patch('llm_cmd.PromptSession') as MockSession:
+        session = MagicMock()
+        MockSession.return_value = session
+        session.prompt.return_value = "test command"
+        yield session
+
+def test_history_saving_with_flag(mock_subprocess_run, mock_subprocess_check_output, mock_prompt_session, monkeypatch, capsys):
+    # Test with --save-history flag
+    monkeypatch.setenv('HISTFILE', '/tmp/test_history')
+    mock_subprocess_run.return_value.returncode = 0
+    mock_subprocess_run.return_value.stderr = ""
+    mock_subprocess_run.return_value.stdout = ""
+
+    mock_subprocess_check_output.return_value = b""
+    interactive_exec("test command", save_history=True)
+
+    # Verify history command was called
+    mock_subprocess_run.assert_called_once()
+    assert 'history -s' in mock_subprocess_run.call_args[0][0][2]
+
+def test_history_saving_with_env_var(mock_subprocess_run, mock_subprocess_check_output, mock_prompt_session, monkeypatch, capsys):
+    # Test with environment variable
+    monkeypatch.setenv('HISTFILE', '/tmp/test_history')
+    monkeypatch.setenv('LLM_CMD_SAVE_HISTORY', '1')
+    mock_subprocess_run.return_value.returncode = 0
+    mock_subprocess_check_output.return_value = b""
+
+    interactive_exec("test command", save_history=False)
+
+    # Verify history command was called due to env var
+    mock_subprocess_run.assert_called_once()
+    assert 'history -s' in mock_subprocess_run.call_args[0][0][2]
+
+def test_no_history_saving(mock_subprocess_run, mock_subprocess_check_output, mock_prompt_session, monkeypatch, capsys):
+    # Test with no history saving enabled
+    monkeypatch.delenv('LLM_CMD_SAVE_HISTORY', raising=False)
+
+    interactive_exec("test command", save_history=False)
+
+    # Verify history command was not called
+    mock_subprocess_run.assert_not_called()
+
+def test_history_saving_failure(mock_subprocess_run, mock_subprocess_check_output, mock_prompt_session, monkeypatch, capsys):
+    # Test history saving with no HISTFILE set
+    monkeypatch.delenv('HISTFILE', raising=False)
+
+    interactive_exec("test command", save_history=True)
+
+    # Verify warning was printed and history command was not called
+    captured = capsys.readouterr()
+    assert "Warning: $HISTFILE environment variable not set or not exported" in captured.out
+    assert "History saving is disabled" in captured.out
+    mock_subprocess_run.assert_not_called()
+
+    # Test history saving failure scenario
+    monkeypatch.setenv('HISTFILE', '/tmp/test_history')
+    mock_subprocess_run.return_value.returncode = 1
+    mock_subprocess_run.return_value.stderr = "Permission denied"
+    mock_subprocess_run.return_value.stdout = ""
+
+    interactive_exec("test command", save_history=True)
+
+    # Verify error messages were printed
+    captured = capsys.readouterr()
+    assert "Warning: Failed to save command to history" in captured.out
+    assert "Permission denied" in captured.out