mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-31 16:21:50 +03:00
chengyongruandGitHub
6cd7063682
* refactor(agent): defer transcript assembly to runner Keep persisted history and the fresh turn as explicit inputs until the Runner assembles the provider transcript. Preserve ContextBuilder and direct AgentRunner compatibility while making the save boundary structural. Refs NAN-81. * fix(providers): preserve mixed adjacent user content
450 lines
17 KiB
Python
450 lines
17 KiB
Python
"""Tests for /restart slash command."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
import sys
|
|
import time
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from nanobot.agent.context import TranscriptInput
|
|
from nanobot.bus.events import InboundMessage
|
|
from nanobot.providers.base import LLMResponse, LLMUsage
|
|
|
|
|
|
def _make_loop():
|
|
"""Create a minimal AgentLoop with mocked dependencies."""
|
|
from nanobot.agent.loop import AgentLoop
|
|
from nanobot.bus.queue import MessageBus
|
|
|
|
bus = MessageBus()
|
|
provider = MagicMock()
|
|
provider.get_default_model.return_value = "test-model"
|
|
workspace = MagicMock()
|
|
workspace.__truediv__ = MagicMock(return_value=MagicMock())
|
|
|
|
with patch("nanobot.agent.loop.ContextBuilder"), \
|
|
patch("nanobot.agent.loop.SessionManager"), \
|
|
patch("nanobot.agent.loop.SubagentManager") as mock_sub_mgr:
|
|
mock_sub_mgr.return_value.close = AsyncMock()
|
|
loop = AgentLoop(bus=bus, provider=provider, workspace=workspace)
|
|
return loop, bus
|
|
|
|
|
|
class TestRestartCommand:
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restart_sends_message_and_calls_execv(self):
|
|
from nanobot.command.builtin import cmd_restart
|
|
from nanobot.command.router import CommandContext
|
|
from nanobot.utils.restart import (
|
|
RESTART_NOTIFY_CHANNEL_ENV,
|
|
RESTART_NOTIFY_CHAT_ID_ENV,
|
|
RESTART_STARTED_AT_ENV,
|
|
)
|
|
|
|
loop, _bus = _make_loop()
|
|
loop.restart_mode = "exec"
|
|
msg = InboundMessage(channel="cli", sender_id="user", chat_id="direct", content="/restart")
|
|
ctx = CommandContext(msg=msg, session=None, key=msg.session_key, raw="/restart", loop=loop)
|
|
|
|
async def _fast_sleep(_delay: float) -> None:
|
|
return None
|
|
|
|
scheduled: list[asyncio.Task] = []
|
|
|
|
def _capture_task(coro):
|
|
task = asyncio.create_task(coro)
|
|
scheduled.append(task)
|
|
return task
|
|
|
|
fake_asyncio = SimpleNamespace(
|
|
sleep=_fast_sleep,
|
|
create_task=_capture_task,
|
|
)
|
|
|
|
with patch.dict(os.environ, {}, clear=False), \
|
|
patch("nanobot.command.builtin.asyncio", new=fake_asyncio), \
|
|
patch("nanobot.command.builtin.os.execv") as mock_execv:
|
|
out = await cmd_restart(ctx)
|
|
assert "Restarting" in out.content
|
|
assert os.environ.get(RESTART_NOTIFY_CHANNEL_ENV) == "cli"
|
|
assert os.environ.get(RESTART_NOTIFY_CHAT_ID_ENV) == "direct"
|
|
assert os.environ.get(RESTART_STARTED_AT_ENV)
|
|
|
|
assert scheduled
|
|
await scheduled[0]
|
|
mock_execv.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restart_windows_auto_spawns_and_exits(self):
|
|
from nanobot.command.builtin import cmd_restart
|
|
from nanobot.command.router import CommandContext
|
|
|
|
loop, _bus = _make_loop()
|
|
msg = InboundMessage(channel="cli", sender_id="user", chat_id="direct", content="/restart")
|
|
ctx = CommandContext(msg=msg, session=None, key=msg.session_key, raw="/restart", loop=loop)
|
|
|
|
async def _fast_sleep(_delay: float) -> None:
|
|
return None
|
|
|
|
scheduled: list[asyncio.Task] = []
|
|
fake_asyncio = SimpleNamespace(
|
|
sleep=_fast_sleep,
|
|
create_task=lambda coro: scheduled.append(asyncio.create_task(coro)) or scheduled[-1],
|
|
)
|
|
|
|
with patch("nanobot.command.builtin.asyncio", new=fake_asyncio), \
|
|
patch("nanobot.command.builtin.sys.platform", "win32"), \
|
|
patch("nanobot.command.builtin.subprocess.CREATE_NEW_PROCESS_GROUP", 512, create=True), \
|
|
patch("nanobot.command.builtin.subprocess.Popen") as mock_popen, \
|
|
patch("nanobot.command.builtin.os._exit") as mock_exit, \
|
|
patch("nanobot.command.builtin.os.execv") as mock_execv:
|
|
await cmd_restart(ctx)
|
|
await scheduled[0]
|
|
|
|
mock_popen.assert_called_once_with(
|
|
[sys.executable, "-m", "nanobot"] + sys.argv[1:],
|
|
creationflags=512,
|
|
)
|
|
mock_exit.assert_called_once_with(0)
|
|
mock_execv.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restart_exit_mode_does_not_spawn(self):
|
|
from nanobot.command.builtin import cmd_restart
|
|
from nanobot.command.router import CommandContext
|
|
|
|
loop, _bus = _make_loop()
|
|
loop.restart_mode = "exit"
|
|
msg = InboundMessage(channel="cli", sender_id="user", chat_id="direct", content="/restart")
|
|
ctx = CommandContext(msg=msg, session=None, key=msg.session_key, raw="/restart", loop=loop)
|
|
|
|
async def _fast_sleep(_delay: float) -> None:
|
|
return None
|
|
|
|
scheduled: list[asyncio.Task] = []
|
|
fake_asyncio = SimpleNamespace(
|
|
sleep=_fast_sleep,
|
|
create_task=lambda coro: scheduled.append(asyncio.create_task(coro)) or scheduled[-1],
|
|
)
|
|
|
|
with patch("nanobot.command.builtin.asyncio", new=fake_asyncio), \
|
|
patch("nanobot.command.builtin.subprocess.Popen") as mock_popen, \
|
|
patch("nanobot.command.builtin.os._exit") as mock_exit, \
|
|
patch("nanobot.command.builtin.os.execv") as mock_execv:
|
|
await cmd_restart(ctx)
|
|
await scheduled[0]
|
|
|
|
mock_exit.assert_called_once_with(0)
|
|
mock_popen.assert_not_called()
|
|
mock_execv.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restart_intercepted_in_run_loop(self):
|
|
"""Verify /restart is handled at the run-loop level, not inside _dispatch."""
|
|
loop, bus = _make_loop()
|
|
loop.restart_mode = "exec"
|
|
msg = InboundMessage(channel="telegram", sender_id="u1", chat_id="c1", content="/restart")
|
|
|
|
async def _fast_sleep(_delay: float) -> None:
|
|
return None
|
|
|
|
scheduled: list[asyncio.Task] = []
|
|
|
|
def _capture_task(coro):
|
|
task = asyncio.create_task(coro)
|
|
scheduled.append(task)
|
|
return task
|
|
|
|
fake_asyncio = SimpleNamespace(
|
|
sleep=_fast_sleep,
|
|
create_task=_capture_task,
|
|
)
|
|
|
|
with patch.object(loop, "_dispatch", new_callable=AsyncMock) as mock_dispatch, \
|
|
patch("nanobot.command.builtin.asyncio", new=fake_asyncio), \
|
|
patch("nanobot.command.builtin.os.execv"):
|
|
await bus.publish_inbound(msg)
|
|
|
|
loop._running = True
|
|
run_task = asyncio.create_task(loop.run())
|
|
out = await asyncio.wait_for(bus.consume_outbound(), timeout=1.0)
|
|
loop._running = False
|
|
run_task.cancel()
|
|
try:
|
|
await run_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
mock_dispatch.assert_not_called()
|
|
assert "Restarting" in out.content
|
|
assert scheduled
|
|
await scheduled[0]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_status_intercepted_in_run_loop(self):
|
|
"""Verify /status is handled at the run-loop level for immediate replies."""
|
|
loop, bus = _make_loop()
|
|
msg = InboundMessage(channel="telegram", sender_id="u1", chat_id="c1", content="/status")
|
|
|
|
with patch.object(loop, "_dispatch", new_callable=AsyncMock) as mock_dispatch:
|
|
await bus.publish_inbound(msg)
|
|
|
|
loop._running = True
|
|
run_task = asyncio.create_task(loop.run())
|
|
out = await asyncio.wait_for(bus.consume_outbound(), timeout=1.0)
|
|
loop._running = False
|
|
run_task.cancel()
|
|
try:
|
|
await run_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
mock_dispatch.assert_not_called()
|
|
assert "nanobot" in out.content.lower() or "Model" in out.content
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_propagates_external_cancellation(self):
|
|
"""External task cancellation should not be swallowed by the inbound wait loop."""
|
|
loop, _bus = _make_loop()
|
|
|
|
run_task = asyncio.create_task(loop.run())
|
|
await asyncio.sleep(0)
|
|
run_task.cancel()
|
|
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await asyncio.wait_for(run_task, timeout=1.0)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_help_includes_restart(self):
|
|
loop, bus = _make_loop()
|
|
msg = InboundMessage(channel="telegram", sender_id="u1", chat_id="c1", content="/help")
|
|
|
|
response = await loop._process_message(msg)
|
|
|
|
assert response is not None
|
|
assert "/restart" in response.content
|
|
assert "/status" in response.content
|
|
assert response.metadata == {"render_as": "text"}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_status_reports_runtime_info(self):
|
|
loop, _bus = _make_loop()
|
|
session = MagicMock()
|
|
session.get_history.return_value = [{"role": "user"}] * 3
|
|
session.metadata = {
|
|
"_last_usage": LLMUsage.reported(input_tokens=0, output_tokens=0).to_dict()
|
|
}
|
|
loop.sessions.get_or_create.return_value = session
|
|
loop._start_time = time.time() - 125
|
|
loop.consolidator.estimate_session_prompt_tokens = MagicMock(
|
|
return_value=(20500, "tiktoken")
|
|
)
|
|
loop.subagents.get_running_count_by_session.return_value = 0
|
|
|
|
msg = InboundMessage(channel="telegram", sender_id="u1", chat_id="c1", content="/status")
|
|
runtime = loop.llm_runtime()
|
|
loop.set_runtime_model("replacement-model")
|
|
loop.set_runtime_context_window(10)
|
|
loop.provider.generation = SimpleNamespace(
|
|
temperature=1.0,
|
|
max_tokens=1,
|
|
reasoning_effort=None,
|
|
)
|
|
|
|
response = await loop._process_message(msg, runtime=runtime)
|
|
|
|
assert response is not None
|
|
assert "Model: test-model" in response.content
|
|
assert "Tokens: 0 in / 0 out" in response.content
|
|
assert "Context: 20k/200k (10% of input budget)" in response.content
|
|
assert "Session: 3 messages" in response.content
|
|
assert "Uptime: 2m 5s" in response.content
|
|
assert "Tasks: 0 active" in response.content
|
|
assert response.metadata == {"render_as": "text"}
|
|
loop.consolidator.estimate_session_prompt_tokens.assert_called_once_with(
|
|
session,
|
|
runtime=runtime,
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_status_counts_running_dispatch_and_subagent_tasks(self):
|
|
loop, _bus = _make_loop()
|
|
session = MagicMock()
|
|
session.get_history.return_value = [{"role": "user"}]
|
|
loop.sessions.get_or_create.return_value = session
|
|
loop.consolidator.estimate_session_prompt_tokens = MagicMock(
|
|
return_value=(1000, "tiktoken")
|
|
)
|
|
|
|
running_task = MagicMock()
|
|
running_task.done.return_value = False
|
|
finished_task = MagicMock()
|
|
finished_task.done.return_value = True
|
|
|
|
msg = InboundMessage(channel="telegram", sender_id="u1", chat_id="c1", content="/status")
|
|
loop._active_tasks[msg.session_key] = {running_task, finished_task}
|
|
loop.subagents.get_running_count_by_session.return_value = 2
|
|
|
|
response = await loop._process_message(msg)
|
|
|
|
assert response is not None
|
|
assert "Tasks: 3 active" in response.content
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_agent_loop_estimates_usage_when_provider_omits_it(self, monkeypatch):
|
|
loop, _bus = _make_loop()
|
|
monkeypatch.setattr(
|
|
"nanobot.agent.runner.estimate_prompt_tokens_chain",
|
|
lambda *_args, **_kwargs: (123, "test"),
|
|
)
|
|
monkeypatch.setattr(
|
|
"nanobot.agent.runner.estimate_message_tokens",
|
|
lambda _message: 7,
|
|
)
|
|
loop.provider.chat_with_retry = AsyncMock(side_effect=[
|
|
LLMResponse(content="first", usage=LLMUsage.reported(input_tokens=9, output_tokens=4)),
|
|
LLMResponse(content="second", usage=None),
|
|
])
|
|
|
|
first = await loop._run_agent_loop(
|
|
TranscriptInput(history=[], current_message=None),
|
|
runtime=loop.llm_runtime(),
|
|
)
|
|
assert first.usage == LLMUsage.reported(input_tokens=9, output_tokens=4)
|
|
|
|
second = await loop._run_agent_loop(
|
|
TranscriptInput(history=[], current_message=None),
|
|
runtime=loop.llm_runtime(),
|
|
)
|
|
assert second.usage == LLMUsage.estimated(input_tokens=123, output_tokens=7)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_status_falls_back_to_session_usage_when_context_estimate_missing(self):
|
|
loop, _bus = _make_loop()
|
|
session = MagicMock()
|
|
session.get_history.return_value = [{"role": "user"}]
|
|
session.metadata = {
|
|
"_last_usage": LLMUsage.reported(input_tokens=1200, output_tokens=34).to_dict()
|
|
}
|
|
loop.sessions.get_or_create.return_value = session
|
|
loop.consolidator.estimate_session_prompt_tokens = MagicMock(
|
|
return_value=(0, "none")
|
|
)
|
|
loop.subagents.get_running_count_by_session.return_value = 0
|
|
|
|
response = await loop._process_message(
|
|
InboundMessage(channel="telegram", sender_id="u1", chat_id="c1", content="/status")
|
|
)
|
|
|
|
assert response is not None
|
|
assert "Tokens: 1200 in / 34 out" in response.content
|
|
assert "Context: 1k/200k (0% of input budget)" in response.content
|
|
assert "Tasks: 0 active" in response.content
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_history_shows_recent_messages(self):
|
|
loop, _bus = _make_loop()
|
|
session = MagicMock()
|
|
session.get_history.return_value = [
|
|
{"role": "user", "content": "Hello"},
|
|
{"role": "assistant", "content": "Hi there!"},
|
|
{"role": "tool", "content": "tool result"}, # should be filtered out
|
|
{"role": "user", "content": "How are you?"},
|
|
{"role": "assistant", "content": "I am doing well."},
|
|
]
|
|
loop.sessions.get_or_create.return_value = session
|
|
|
|
msg = InboundMessage(channel="telegram", sender_id="u1", chat_id="c1", content="/history")
|
|
response = await loop._process_message(msg)
|
|
|
|
assert response is not None
|
|
assert "👤 You: Hello" in response.content
|
|
assert "🤖 Bot: Hi there!" in response.content
|
|
assert "tool result" not in response.content # tool messages filtered
|
|
assert response.metadata == {"render_as": "text"}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_history_respects_count_argument(self):
|
|
loop, _bus = _make_loop()
|
|
session = MagicMock()
|
|
session.get_history.return_value = [
|
|
{"role": "user", "content": f"message {i}"} for i in range(20)
|
|
]
|
|
loop.sessions.get_or_create.return_value = session
|
|
|
|
msg = InboundMessage(channel="telegram", sender_id="u1", chat_id="c1", content="/history 3")
|
|
response = await loop._process_message(msg)
|
|
|
|
assert response is not None
|
|
assert "Last 3 message(s)" in response.content
|
|
assert "message 19" in response.content # most recent
|
|
assert "message 0" not in response.content # too old
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_history_clamps_count_and_extracts_text_blocks(self):
|
|
loop, _bus = _make_loop()
|
|
session = MagicMock()
|
|
session.get_history.return_value = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "visible text"},
|
|
{"type": "image_url", "image_url": {"url": "data:image/png;base64,..."}},
|
|
],
|
|
},
|
|
*({"role": "assistant", "content": f"reply {i}"} for i in range(60)),
|
|
]
|
|
loop.sessions.get_or_create.return_value = session
|
|
|
|
msg = InboundMessage(channel="telegram", sender_id="u1", chat_id="c1", content="/history 999")
|
|
response = await loop._process_message(msg)
|
|
|
|
assert response is not None
|
|
assert "Last 50 message(s)" in response.content
|
|
assert "visible text" not in response.content
|
|
assert "reply 59" in response.content
|
|
assert "reply 9" not in response.content
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_history_invalid_count_returns_usage(self):
|
|
loop, _bus = _make_loop()
|
|
|
|
msg = InboundMessage(channel="telegram", sender_id="u1", chat_id="c1", content="/history nope")
|
|
response = await loop._process_message(msg)
|
|
|
|
assert response is not None
|
|
assert response.content.startswith("Usage: /history [count]")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_history_empty_session(self):
|
|
loop, _bus = _make_loop()
|
|
session = MagicMock()
|
|
session.get_history.return_value = []
|
|
loop.sessions.get_or_create.return_value = session
|
|
|
|
msg = InboundMessage(channel="telegram", sender_id="u1", chat_id="c1", content="/history")
|
|
response = await loop._process_message(msg)
|
|
|
|
assert response is not None
|
|
assert "No conversation history yet." in response.content
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_direct_preserves_render_metadata(self):
|
|
loop, _bus = _make_loop()
|
|
session = MagicMock()
|
|
session.get_history.return_value = []
|
|
loop.sessions.get_or_create.return_value = session
|
|
loop.subagents.get_running_count.return_value = 0
|
|
loop.subagents.get_running_count_by_session.return_value = 0
|
|
|
|
response = await loop.process_direct("/status", session_key="cli:test")
|
|
|
|
assert response is not None
|
|
assert response.metadata == {"render_as": "text"}
|