forked from Rockachopa/Timmy-time-dashboard
Kimi Agent
f8f3b9b81f
Timmy can now introspect which session he's running in (cli, dashboard, loop).
- Add {session_id} placeholder to both lite and full system prompts
- get_system_prompt() accepts session_id param (default: 'unknown')
- create_timmy() accepts session_id param, forwards to prompt
- CLI chat/think/status pass their session_id to create_timmy()
- session.py passes _DEFAULT_SESSION_ID to create_timmy()
- 7 new tests in test_session_identity.py
- Updated 2 existing CLI test mocks
Closes #64
280 lines
9.3 KiB
Python
280 lines
9.3 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_passes_model_size_option():
|
|
"""think --model-size 70b must forward the 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, model_size="70b", 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():
|
|
"""chat --backend airllm must forward the backend to create_timmy."""
|
|
mock_run_output = MagicMock()
|
|
mock_run_output.content = "OK"
|
|
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) as mock_create:
|
|
runner.invoke(app, ["chat", "test", "--backend", "airllm"])
|
|
|
|
mock_create.assert_called_once_with(backend="airllm", model_size=None, session_id="cli")
|
|
|
|
|
|
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
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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()
|