diff --git a/shellmind/ai.py b/shellmind/ai.py index bb4cbdb..3530361 100644 --- a/shellmind/ai.py +++ b/shellmind/ai.py @@ -211,11 +211,15 @@ def _not_set_error(name: str) -> None: banner(console) - try: - while True: - uin = input(config.prompt) + while True: + try: + try: + uin = input(config.prompt) + except EOFError: + print() + return - if uin == "/bye": + if uin in ("/bye", "exit", "quit"): return if uin == "/model": @@ -256,7 +260,7 @@ def _not_set_error(name: str) -> None: {Fore.CYAN}Console{Fore.YELLOW} Type {Fore.BLUE}/model{Fore.YELLOW} to show active model. Type {Fore.BLUE}/help{Fore.YELLOW} or {Fore.BLUE}/?{Fore.YELLOW} to see this. -Type {Fore.BLUE}/bye{Fore.YELLOW} to exit. +Type {Fore.BLUE}/bye{Fore.YELLOW}, {Fore.BLUE}exit{Fore.YELLOW} or {Fore.BLUE}quit{Fore.YELLOW} to exit. Type {Fore.BLUE}/clear{Fore.YELLOW} to clear saved context. Type {Fore.BLUE}!{Fore.YELLOW} to run shell directly. Type {Fore.BLUE}/run {Fore.YELLOW} @@ -377,8 +381,9 @@ def _not_set_error(name: str) -> None: print() - except KeyboardInterrupt: - return + except KeyboardInterrupt: + print() + continue if __name__ == "__main__": diff --git a/shellmind/tools.py b/shellmind/tools.py index c34d170..96c9407 100644 --- a/shellmind/tools.py +++ b/shellmind/tools.py @@ -12,28 +12,34 @@ def shell_tool(command: str) -> str: print(Fore.BLUE + f"Executing shell command: {command}" + Style.RESET_ALL) - if os.name == "nt": - args = [ - "powershell", - "-NoProfile", - "-ExecutionPolicy", - "Bypass", - "-Command", - command, - ] - result = subprocess.run( - args, - capture_output=True, - text=True, - timeout=15, - ) - else: - result = subprocess.run( - command, - shell=True, - capture_output=True, - text=True, - timeout=15, + try: + if os.name == "nt": + args = [ + "powershell", + "-NoProfile", + "-ExecutionPolicy", + "Bypass", + "-Command", + command, + ] + result = subprocess.run( + args, + capture_output=True, + text=True, + timeout=15, + ) + else: + result = subprocess.run( + command, + shell=True, + capture_output=True, + text=True, + timeout=15, + ) + except subprocess.TimeoutExpired: + return ( + "Error: Command timed out after 15 seconds. " + "Please note that shell commands are run non-interactively." ) parts = [result.stdout.strip(), result.stderr.strip()] diff --git a/tests/test_tools.py b/tests/test_tools.py new file mode 100644 index 0000000..646ecd4 --- /dev/null +++ b/tests/test_tools.py @@ -0,0 +1,28 @@ +import subprocess +from shellmind.tools import shell_tool + +def test_shell_tool_timeout(monkeypatch): + # This might be tricky to test without actually waiting 15 seconds + # but we can mock subprocess.run to raise TimeoutExpired + def mock_run(*args, **kwargs): + raise subprocess.TimeoutExpired(args[0], 15) + + monkeypatch.setattr(subprocess, "run", mock_run) + + result = shell_tool("long_command") + assert "Command timed out" in result + assert "non-interactively" in result + +def test_shell_tool_success(monkeypatch): + class MockResult: + stdout = "success" + stderr = "" + returncode = 0 + + def mock_run(*args, **kwargs): + return MockResult() + + monkeypatch.setattr(subprocess, "run", mock_run) + + result = shell_tool("echo success") + assert result == "success"