Kimi Agent
e3d60b89a9
Co-authored-by: Kimi Agent <kimi@timmy.local> Co-committed-by: Kimi Agent <kimi@timmy.local>
442 lines
14 KiB
Python
442 lines
14 KiB
Python
from unittest.mock import MagicMock, patch
|
|
|
|
from typer.testing import CliRunner
|
|
|
|
from timmy.cli import _CLI_SESSION_ID, _handle_tool_confirmation, app
|
|
from timmy.prompts import STATUS_PROMPT
|
|
|
|
runner = CliRunner()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# status command
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_status_uses_status_prompt():
|
|
"""status command must pass STATUS_PROMPT to the agent."""
|
|
mock_timmy = MagicMock()
|
|
|
|
with patch("timmy.cli.create_timmy", return_value=mock_timmy):
|
|
runner.invoke(app, ["status"])
|
|
|
|
mock_timmy.print_response.assert_called_once_with(
|
|
STATUS_PROMPT, stream=False, session_id=_CLI_SESSION_ID
|
|
)
|
|
|
|
|
|
def test_status_does_not_use_inline_string():
|
|
"""status command must not pass the old inline hardcoded string."""
|
|
mock_timmy = MagicMock()
|
|
|
|
with patch("timmy.cli.create_timmy", return_value=mock_timmy):
|
|
runner.invoke(app, ["status"])
|
|
|
|
call_args = mock_timmy.print_response.call_args
|
|
assert call_args[0][0] != "Brief status report — one sentence."
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# think command
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_think_sends_topic_to_agent():
|
|
"""think command must pass the topic wrapped in a prompt with streaming."""
|
|
mock_timmy = MagicMock()
|
|
|
|
with patch("timmy.cli.create_timmy", return_value=mock_timmy):
|
|
runner.invoke(app, ["think", "Bitcoin self-custody"])
|
|
|
|
mock_timmy.print_response.assert_called_once_with(
|
|
"Think carefully about: Bitcoin self-custody",
|
|
stream=True,
|
|
session_id=_CLI_SESSION_ID,
|
|
)
|
|
|
|
|
|
def test_think_ignores_model_size_option():
|
|
"""think --model-size 70b must not forward model_size to create_timmy."""
|
|
mock_timmy = MagicMock()
|
|
|
|
with patch("timmy.cli.create_timmy", return_value=mock_timmy) as mock_create:
|
|
runner.invoke(app, ["think", "topic", "--model-size", "70b"])
|
|
|
|
mock_create.assert_called_once_with(backend=None, session_id="cli")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# chat command — session persistence
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_chat_uses_session_id():
|
|
"""chat command must pass the stable CLI session_id to agent.run()."""
|
|
mock_run_output = MagicMock()
|
|
mock_run_output.content = "Hello there!"
|
|
mock_run_output.status = "COMPLETED"
|
|
mock_run_output.active_requirements = []
|
|
|
|
mock_timmy = MagicMock()
|
|
mock_timmy.run.return_value = mock_run_output
|
|
|
|
with patch("timmy.cli.create_timmy", return_value=mock_timmy):
|
|
result = runner.invoke(app, ["chat", "Hello Timmy"])
|
|
|
|
mock_timmy.run.assert_called_once_with("Hello Timmy", stream=False, session_id=_CLI_SESSION_ID)
|
|
assert result.exit_code == 0
|
|
assert "Hello there!" in result.output
|
|
|
|
|
|
def test_chat_new_session_uses_unique_id():
|
|
"""chat --new must use a unique session_id, not the stable one."""
|
|
mock_run_output = MagicMock()
|
|
mock_run_output.content = "Fresh start!"
|
|
mock_run_output.status = "COMPLETED"
|
|
mock_run_output.active_requirements = []
|
|
|
|
mock_timmy = MagicMock()
|
|
mock_timmy.run.return_value = mock_run_output
|
|
|
|
with patch("timmy.cli.create_timmy", return_value=mock_timmy):
|
|
runner.invoke(app, ["chat", "Hello", "--new"])
|
|
|
|
call_args = mock_timmy.run.call_args
|
|
used_session_id = call_args[1]["session_id"]
|
|
assert used_session_id != _CLI_SESSION_ID # Must be unique
|
|
|
|
|
|
def test_chat_passes_backend_option():
|
|
pass
|
|
|
|
|
|
def test_chat_cleans_response():
|
|
"""chat must clean tool-call artifacts from the response."""
|
|
raw = '{"name": "python", "parameters": {"code": "1+1"}} The answer is 2.'
|
|
mock_run_output = MagicMock()
|
|
mock_run_output.content = raw
|
|
mock_run_output.status = "COMPLETED"
|
|
mock_run_output.active_requirements = []
|
|
|
|
mock_timmy = MagicMock()
|
|
mock_timmy.run.return_value = mock_run_output
|
|
|
|
with patch("timmy.cli.create_timmy", return_value=mock_timmy):
|
|
result = runner.invoke(app, ["chat", "what is 1+1"])
|
|
|
|
# The JSON tool call should be stripped
|
|
assert '"name": "python"' not in result.output
|
|
assert "The answer is 2." in result.output
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# chat command — multi-word messages
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_chat_multiword_without_quotes():
|
|
"""chat should accept multiple arguments without requiring quotes."""
|
|
mock_run_output = MagicMock()
|
|
mock_run_output.content = "Response"
|
|
mock_run_output.status = "COMPLETED"
|
|
mock_run_output.active_requirements = []
|
|
|
|
mock_timmy = MagicMock()
|
|
mock_timmy.run.return_value = mock_run_output
|
|
|
|
with patch("timmy.cli.create_timmy", return_value=mock_timmy):
|
|
result = runner.invoke(app, ["chat", "hello", "world", "how", "are", "you"])
|
|
|
|
mock_timmy.run.assert_called_once_with(
|
|
"hello world how are you", stream=False, session_id=_CLI_SESSION_ID
|
|
)
|
|
assert result.exit_code == 0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# chat command — stdin support
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_chat_stdin_with_dash():
|
|
"""chat - should read message from stdin when "-" is passed."""
|
|
mock_run_output = MagicMock()
|
|
mock_run_output.content = "Stdin response"
|
|
mock_run_output.status = "COMPLETED"
|
|
mock_run_output.active_requirements = []
|
|
|
|
mock_timmy = MagicMock()
|
|
mock_timmy.run.return_value = mock_run_output
|
|
|
|
with patch("timmy.cli.create_timmy", return_value=mock_timmy):
|
|
result = runner.invoke(app, ["chat", "-"], input="Hello from stdin")
|
|
|
|
mock_timmy.run.assert_called_once_with(
|
|
"Hello from stdin", stream=False, session_id=_CLI_SESSION_ID
|
|
)
|
|
assert result.exit_code == 0
|
|
assert "Stdin response" in result.output
|
|
|
|
|
|
def test_chat_stdin_piped():
|
|
"""chat should read from stdin when piped (non-tty)."""
|
|
mock_run_output = MagicMock()
|
|
mock_run_output.content = "Piped response"
|
|
mock_run_output.status = "COMPLETED"
|
|
mock_run_output.active_requirements = []
|
|
|
|
mock_timmy = MagicMock()
|
|
mock_timmy.run.return_value = mock_run_output
|
|
|
|
with (
|
|
patch("timmy.cli.create_timmy", return_value=mock_timmy),
|
|
patch("timmy.cli._is_interactive", return_value=False),
|
|
):
|
|
result = runner.invoke(app, ["chat", "ignored"], input="Piped message")
|
|
|
|
mock_timmy.run.assert_called_once_with(
|
|
"Piped message", stream=False, session_id=_CLI_SESSION_ID
|
|
)
|
|
assert result.exit_code == 0
|
|
|
|
|
|
def test_chat_stdin_empty_exits_with_error():
|
|
"""chat - should exit with error when stdin is empty."""
|
|
with patch("timmy.cli._is_interactive", return_value=False):
|
|
result = runner.invoke(app, ["chat", "-"], input="")
|
|
|
|
assert result.exit_code == 1
|
|
assert "No input provided" in result.output or "No input provided" in str(result.output)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# repl command
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_repl_exits_on_ctrl_c():
|
|
"""repl should exit gracefully on Ctrl+C (KeyboardInterrupt)."""
|
|
with patch("builtins.input", side_effect=KeyboardInterrupt):
|
|
result = runner.invoke(app, ["repl"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "Goodbye" in result.output
|
|
|
|
|
|
def test_repl_exits_on_ctrl_d():
|
|
"""repl should exit gracefully on Ctrl+D (EOFError)."""
|
|
with patch("builtins.input", side_effect=EOFError):
|
|
result = runner.invoke(app, ["repl"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "Goodbye" in result.output
|
|
|
|
|
|
def test_repl_exits_on_quit_command():
|
|
"""repl should exit when user types 'quit'."""
|
|
with patch("builtins.input", return_value="quit"):
|
|
result = runner.invoke(app, ["repl"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "Goodbye" in result.output
|
|
|
|
|
|
def test_repl_exits_on_exit_command():
|
|
"""repl should exit when user types 'exit'."""
|
|
with patch("builtins.input", return_value="exit"):
|
|
result = runner.invoke(app, ["repl"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "Goodbye" in result.output
|
|
|
|
|
|
def test_repl_sends_message_to_chat():
|
|
"""repl should send user input to session.chat()."""
|
|
with (
|
|
patch("builtins.input", side_effect=["Hello Timmy", "exit"]),
|
|
patch("timmy.session.chat") as mock_chat,
|
|
):
|
|
mock_chat.return_value = "Hi there!"
|
|
result = runner.invoke(app, ["repl"])
|
|
|
|
mock_chat.assert_called()
|
|
assert result.exit_code == 0
|
|
|
|
|
|
def test_repl_skips_empty_input():
|
|
"""repl should skip empty input lines."""
|
|
with (
|
|
patch("builtins.input", side_effect=["", "", "hello", "exit"]),
|
|
patch("timmy.session.chat") as mock_chat,
|
|
):
|
|
mock_chat.return_value = "Response"
|
|
runner.invoke(app, ["repl"])
|
|
|
|
# chat should only be called once (for "hello"), empty lines are skipped, exit breaks
|
|
assert mock_chat.call_count == 1
|
|
|
|
|
|
def test_repl_uses_custom_session_id():
|
|
"""repl --session-id should use the provided session ID."""
|
|
with (
|
|
patch("builtins.input", side_effect=["hello", "exit"]),
|
|
patch("timmy.session.chat") as mock_chat,
|
|
):
|
|
mock_chat.return_value = "Response"
|
|
runner.invoke(app, ["repl", "--session-id", "custom-session"])
|
|
|
|
# Check that chat was called with the custom session_id
|
|
calls = mock_chat.call_args_list
|
|
assert any(call.kwargs.get("session_id") == "custom-session" for call in calls)
|
|
|
|
|
|
def test_repl_handles_chat_errors():
|
|
"""repl should handle errors from chat() gracefully."""
|
|
with (
|
|
patch("builtins.input", side_effect=["hello", "exit"]),
|
|
patch("timmy.session.chat") as mock_chat,
|
|
):
|
|
mock_chat.side_effect = Exception("Chat error")
|
|
result = runner.invoke(app, ["repl"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "Error" in result.output
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tool confirmation gate
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_paused_run(tool_name="shell", tool_args=None):
|
|
"""Create a mock paused RunOutput with one requirement."""
|
|
tool_args = tool_args or {"command": "ls -la"}
|
|
|
|
mock_te = MagicMock()
|
|
mock_te.tool_name = tool_name
|
|
mock_te.tool_args = tool_args
|
|
|
|
mock_req = MagicMock()
|
|
mock_req.needs_confirmation = True
|
|
mock_req.tool_execution = mock_te
|
|
|
|
mock_run = MagicMock()
|
|
mock_run.status = "RunStatus.paused"
|
|
mock_run.active_requirements = [mock_req]
|
|
|
|
return mock_run, mock_req
|
|
|
|
|
|
def test_handle_tool_confirmation_approve():
|
|
"""Approving a tool should call req.confirm() and agent.continue_run()."""
|
|
paused_run, mock_req = _make_paused_run()
|
|
|
|
completed_run = MagicMock()
|
|
completed_run.status = "COMPLETED"
|
|
completed_run.active_requirements = []
|
|
completed_run.content = "Done."
|
|
|
|
mock_agent = MagicMock()
|
|
mock_agent.continue_run.return_value = completed_run
|
|
|
|
# Simulate user typing "y" at the prompt (mock interactive terminal)
|
|
with (
|
|
patch("timmy.cli._is_interactive", return_value=True),
|
|
patch("timmy.cli.typer.confirm", return_value=True),
|
|
):
|
|
result = _handle_tool_confirmation(mock_agent, paused_run, "cli")
|
|
|
|
mock_req.confirm.assert_called_once()
|
|
mock_agent.continue_run.assert_called_once()
|
|
assert result.content == "Done."
|
|
|
|
|
|
def test_handle_tool_confirmation_reject():
|
|
"""Rejecting a tool should call req.reject() and agent.continue_run()."""
|
|
paused_run, mock_req = _make_paused_run()
|
|
|
|
completed_run = MagicMock()
|
|
completed_run.status = "COMPLETED"
|
|
completed_run.active_requirements = []
|
|
completed_run.content = "Action rejected."
|
|
|
|
mock_agent = MagicMock()
|
|
mock_agent.continue_run.return_value = completed_run
|
|
|
|
with (
|
|
patch("timmy.cli._is_interactive", return_value=True),
|
|
patch("timmy.cli.typer.confirm", return_value=False),
|
|
):
|
|
_handle_tool_confirmation(mock_agent, paused_run, "cli")
|
|
|
|
mock_req.reject.assert_called_once()
|
|
mock_agent.continue_run.assert_called_once()
|
|
|
|
|
|
def test_handle_tool_confirmation_not_paused():
|
|
"""Non-paused runs should pass through unchanged."""
|
|
completed_run = MagicMock()
|
|
completed_run.status = "COMPLETED"
|
|
completed_run.active_requirements = []
|
|
|
|
mock_agent = MagicMock()
|
|
result = _handle_tool_confirmation(mock_agent, completed_run, "cli")
|
|
|
|
assert result is completed_run
|
|
mock_agent.continue_run.assert_not_called()
|
|
|
|
|
|
def test_handle_tool_confirmation_continue_error():
|
|
"""Errors in continue_run should be handled gracefully."""
|
|
paused_run, mock_req = _make_paused_run()
|
|
|
|
mock_agent = MagicMock()
|
|
mock_agent.continue_run.side_effect = Exception("connection lost")
|
|
|
|
with (
|
|
patch("timmy.cli._is_interactive", return_value=True),
|
|
patch("timmy.cli.typer.confirm", return_value=True),
|
|
):
|
|
result = _handle_tool_confirmation(mock_agent, paused_run, "cli")
|
|
|
|
# Should return the original paused run, not crash
|
|
assert result is paused_run
|
|
|
|
|
|
def test_handle_tool_confirmation_autonomous_allowlisted():
|
|
"""In autonomous mode, allowlisted tools should be auto-approved."""
|
|
paused_run, mock_req = _make_paused_run(
|
|
tool_name="shell", tool_args={"command": "pytest tests/ -x"}
|
|
)
|
|
|
|
completed_run = MagicMock()
|
|
completed_run.status = "COMPLETED"
|
|
completed_run.active_requirements = []
|
|
|
|
mock_agent = MagicMock()
|
|
mock_agent.continue_run.return_value = completed_run
|
|
|
|
with patch("timmy.cli.is_allowlisted", return_value=True):
|
|
_handle_tool_confirmation(mock_agent, paused_run, "cli", autonomous=True)
|
|
|
|
mock_req.confirm.assert_called_once()
|
|
mock_req.reject.assert_not_called()
|
|
|
|
|
|
def test_handle_tool_confirmation_autonomous_not_allowlisted():
|
|
"""In autonomous mode, non-allowlisted tools should be auto-rejected."""
|
|
paused_run, mock_req = _make_paused_run(tool_name="shell", tool_args={"command": "rm -rf /"})
|
|
|
|
completed_run = MagicMock()
|
|
completed_run.status = "COMPLETED"
|
|
completed_run.active_requirements = []
|
|
|
|
mock_agent = MagicMock()
|
|
mock_agent.continue_run.return_value = completed_run
|
|
|
|
with patch("timmy.cli.is_allowlisted", return_value=False):
|
|
_handle_tool_confirmation(mock_agent, paused_run, "cli", autonomous=True)
|
|
|
|
mock_req.reject.assert_called_once()
|
|
mock_req.confirm.assert_not_called()
|