Merge remote-tracking branch 'origin/main'
This commit is contained in:
@@ -14,19 +14,31 @@ class _FakeResponse:
|
||||
self.status_code = status_code
|
||||
self._json_body = json_body or {}
|
||||
self.text = "{}"
|
||||
self.content = b""
|
||||
self.headers = {"content-type": "application/json"}
|
||||
|
||||
def json(self) -> dict:
|
||||
return self._json_body
|
||||
|
||||
|
||||
class _FakeHttp:
|
||||
def __init__(self) -> None:
|
||||
def __init__(self, responses: list[_FakeResponse] | None = None) -> None:
|
||||
self.calls: list[dict] = []
|
||||
self._responses = list(responses) if responses else []
|
||||
|
||||
async def post(self, url: str, json=None, headers=None):
|
||||
self.calls.append({"url": url, "json": json, "headers": headers})
|
||||
def _next_response(self) -> _FakeResponse:
|
||||
if self._responses:
|
||||
return self._responses.pop(0)
|
||||
return _FakeResponse()
|
||||
|
||||
async def post(self, url: str, json=None, headers=None, **kwargs):
|
||||
self.calls.append({"method": "POST", "url": url, "json": json, "headers": headers})
|
||||
return self._next_response()
|
||||
|
||||
async def get(self, url: str, **kwargs):
|
||||
self.calls.append({"method": "GET", "url": url})
|
||||
return self._next_response()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_group_message_keeps_sender_id_and_routes_chat_id() -> None:
|
||||
@@ -109,3 +121,93 @@ async def test_handler_uses_voice_recognition_text_when_text_is_empty(monkeypatc
|
||||
assert msg.content == "voice transcript"
|
||||
assert msg.sender_id == "user1"
|
||||
assert msg.chat_id == "group:conv123"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handler_processes_file_message(monkeypatch) -> None:
|
||||
"""Test that file messages are handled and forwarded with downloaded path."""
|
||||
bus = MessageBus()
|
||||
channel = DingTalkChannel(
|
||||
DingTalkConfig(client_id="app", client_secret="secret", allow_from=["user1"]),
|
||||
bus,
|
||||
)
|
||||
handler = NanobotDingTalkHandler(channel)
|
||||
|
||||
class _FakeFileChatbotMessage:
|
||||
text = None
|
||||
extensions = {}
|
||||
image_content = None
|
||||
rich_text_content = None
|
||||
sender_staff_id = "user1"
|
||||
sender_id = "fallback-user"
|
||||
sender_nick = "Alice"
|
||||
message_type = "file"
|
||||
|
||||
@staticmethod
|
||||
def from_dict(_data):
|
||||
return _FakeFileChatbotMessage()
|
||||
|
||||
async def fake_download(download_code, filename, sender_id):
|
||||
return f"/tmp/nanobot_dingtalk/{sender_id}/{filename}"
|
||||
|
||||
monkeypatch.setattr(dingtalk_module, "ChatbotMessage", _FakeFileChatbotMessage)
|
||||
monkeypatch.setattr(dingtalk_module, "AckMessage", SimpleNamespace(STATUS_OK="OK"))
|
||||
monkeypatch.setattr(channel, "_download_dingtalk_file", fake_download)
|
||||
|
||||
status, body = await handler.process(
|
||||
SimpleNamespace(
|
||||
data={
|
||||
"conversationType": "1",
|
||||
"content": {"downloadCode": "abc123", "fileName": "report.xlsx"},
|
||||
"text": {"content": ""},
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
await asyncio.gather(*list(channel._background_tasks))
|
||||
msg = await bus.consume_inbound()
|
||||
|
||||
assert (status, body) == ("OK", "OK")
|
||||
assert "[File]" in msg.content
|
||||
assert "/tmp/nanobot_dingtalk/user1/report.xlsx" in msg.content
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_dingtalk_file(tmp_path, monkeypatch) -> None:
|
||||
"""Test the two-step file download flow (get URL then download content)."""
|
||||
channel = DingTalkChannel(
|
||||
DingTalkConfig(client_id="app", client_secret="secret", allow_from=["*"]),
|
||||
MessageBus(),
|
||||
)
|
||||
|
||||
# Mock access token
|
||||
async def fake_get_token():
|
||||
return "test-token"
|
||||
|
||||
monkeypatch.setattr(channel, "_get_access_token", fake_get_token)
|
||||
|
||||
# Mock HTTP: first POST returns downloadUrl, then GET returns file bytes
|
||||
file_content = b"fake file content"
|
||||
channel._http = _FakeHttp(responses=[
|
||||
_FakeResponse(200, {"downloadUrl": "https://example.com/tmpfile"}),
|
||||
_FakeResponse(200),
|
||||
])
|
||||
channel._http._responses[1].content = file_content
|
||||
|
||||
# Redirect media dir to tmp_path
|
||||
monkeypatch.setattr(
|
||||
"nanobot.config.paths.get_media_dir",
|
||||
lambda channel_name=None: tmp_path / channel_name if channel_name else tmp_path,
|
||||
)
|
||||
|
||||
result = await channel._download_dingtalk_file("code123", "test.xlsx", "user1")
|
||||
|
||||
assert result is not None
|
||||
assert result.endswith("test.xlsx")
|
||||
assert (tmp_path / "dingtalk" / "user1" / "test.xlsx").read_bytes() == file_content
|
||||
|
||||
# Verify API calls
|
||||
assert channel._http.calls[0]["method"] == "POST"
|
||||
assert "messageFiles/download" in channel._http.calls[0]["url"]
|
||||
assert channel._http.calls[0]["json"]["downloadCode"] == "code123"
|
||||
assert channel._http.calls[1]["method"] == "GET"
|
||||
|
||||
@@ -249,3 +249,114 @@ class TestListDirTool:
|
||||
result = await tool.execute(path=str(tmp_path / "nope"))
|
||||
assert "Error" in result
|
||||
assert "not found" in result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Workspace restriction + extra_allowed_dirs
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestWorkspaceRestriction:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_blocked_outside_workspace(self, tmp_path):
|
||||
workspace = tmp_path / "ws"
|
||||
workspace.mkdir()
|
||||
outside = tmp_path / "outside"
|
||||
outside.mkdir()
|
||||
secret = outside / "secret.txt"
|
||||
secret.write_text("top secret")
|
||||
|
||||
tool = ReadFileTool(workspace=workspace, allowed_dir=workspace)
|
||||
result = await tool.execute(path=str(secret))
|
||||
assert "Error" in result
|
||||
assert "outside" in result.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_allowed_with_extra_dir(self, tmp_path):
|
||||
workspace = tmp_path / "ws"
|
||||
workspace.mkdir()
|
||||
skills_dir = tmp_path / "skills"
|
||||
skills_dir.mkdir()
|
||||
skill_file = skills_dir / "test_skill" / "SKILL.md"
|
||||
skill_file.parent.mkdir()
|
||||
skill_file.write_text("# Test Skill\nDo something.")
|
||||
|
||||
tool = ReadFileTool(
|
||||
workspace=workspace, allowed_dir=workspace,
|
||||
extra_allowed_dirs=[skills_dir],
|
||||
)
|
||||
result = await tool.execute(path=str(skill_file))
|
||||
assert "Test Skill" in result
|
||||
assert "Error" not in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extra_dirs_does_not_widen_write(self, tmp_path):
|
||||
from nanobot.agent.tools.filesystem import WriteFileTool
|
||||
|
||||
workspace = tmp_path / "ws"
|
||||
workspace.mkdir()
|
||||
outside = tmp_path / "outside"
|
||||
outside.mkdir()
|
||||
|
||||
tool = WriteFileTool(workspace=workspace, allowed_dir=workspace)
|
||||
result = await tool.execute(path=str(outside / "hack.txt"), content="pwned")
|
||||
assert "Error" in result
|
||||
assert "outside" in result.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_still_blocked_for_unrelated_dir(self, tmp_path):
|
||||
workspace = tmp_path / "ws"
|
||||
workspace.mkdir()
|
||||
skills_dir = tmp_path / "skills"
|
||||
skills_dir.mkdir()
|
||||
unrelated = tmp_path / "other"
|
||||
unrelated.mkdir()
|
||||
secret = unrelated / "secret.txt"
|
||||
secret.write_text("nope")
|
||||
|
||||
tool = ReadFileTool(
|
||||
workspace=workspace, allowed_dir=workspace,
|
||||
extra_allowed_dirs=[skills_dir],
|
||||
)
|
||||
result = await tool.execute(path=str(secret))
|
||||
assert "Error" in result
|
||||
assert "outside" in result.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_workspace_file_still_readable_with_extra_dirs(self, tmp_path):
|
||||
"""Adding extra_allowed_dirs must not break normal workspace reads."""
|
||||
workspace = tmp_path / "ws"
|
||||
workspace.mkdir()
|
||||
ws_file = workspace / "README.md"
|
||||
ws_file.write_text("hello from workspace")
|
||||
skills_dir = tmp_path / "skills"
|
||||
skills_dir.mkdir()
|
||||
|
||||
tool = ReadFileTool(
|
||||
workspace=workspace, allowed_dir=workspace,
|
||||
extra_allowed_dirs=[skills_dir],
|
||||
)
|
||||
result = await tool.execute(path=str(ws_file))
|
||||
assert "hello from workspace" in result
|
||||
assert "Error" not in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_blocked_in_extra_dir(self, tmp_path):
|
||||
"""edit_file must not be able to modify files in extra_allowed_dirs."""
|
||||
workspace = tmp_path / "ws"
|
||||
workspace.mkdir()
|
||||
skills_dir = tmp_path / "skills"
|
||||
skills_dir.mkdir()
|
||||
skill_file = skills_dir / "weather" / "SKILL.md"
|
||||
skill_file.parent.mkdir()
|
||||
skill_file.write_text("# Weather\nOriginal content.")
|
||||
|
||||
tool = EditFileTool(workspace=workspace, allowed_dir=workspace)
|
||||
result = await tool.execute(
|
||||
path=str(skill_file),
|
||||
old_text="Original content.",
|
||||
new_text="Hacked content.",
|
||||
)
|
||||
assert "Error" in result
|
||||
assert "outside" in result.lower()
|
||||
assert skill_file.read_text() == "# Weather\nOriginal content."
|
||||
|
||||
@@ -123,3 +123,87 @@ async def test_chat_with_retry_explicit_override_beats_defaults() -> None:
|
||||
assert provider.last_kwargs["temperature"] == 0.9
|
||||
assert provider.last_kwargs["max_tokens"] == 9999
|
||||
assert provider.last_kwargs["reasoning_effort"] == "low"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Image-unsupported fallback tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_IMAGE_MSG = [
|
||||
{"role": "user", "content": [
|
||||
{"type": "text", "text": "describe this"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}},
|
||||
]},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_unsupported_error_retries_without_images() -> None:
|
||||
"""If the model rejects image_url, retry once with images stripped."""
|
||||
provider = ScriptedProvider([
|
||||
LLMResponse(
|
||||
content="Invalid content type. image_url is only supported by certain models",
|
||||
finish_reason="error",
|
||||
),
|
||||
LLMResponse(content="ok, no image"),
|
||||
])
|
||||
|
||||
response = await provider.chat_with_retry(messages=_IMAGE_MSG)
|
||||
|
||||
assert response.content == "ok, no image"
|
||||
assert provider.calls == 2
|
||||
msgs_on_retry = provider.last_kwargs["messages"]
|
||||
for msg in msgs_on_retry:
|
||||
content = msg.get("content")
|
||||
if isinstance(content, list):
|
||||
assert all(b.get("type") != "image_url" for b in content)
|
||||
assert any("[image omitted]" in (b.get("text") or "") for b in content)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_unsupported_error_no_retry_without_image_content() -> None:
|
||||
"""If messages don't contain image_url blocks, don't retry on image error."""
|
||||
provider = ScriptedProvider([
|
||||
LLMResponse(
|
||||
content="image_url is only supported by certain models",
|
||||
finish_reason="error",
|
||||
),
|
||||
])
|
||||
|
||||
response = await provider.chat_with_retry(
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
)
|
||||
|
||||
assert provider.calls == 1
|
||||
assert response.finish_reason == "error"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_unsupported_fallback_returns_error_on_second_failure() -> None:
|
||||
"""If the image-stripped retry also fails, return that error."""
|
||||
provider = ScriptedProvider([
|
||||
LLMResponse(
|
||||
content="does not support image input",
|
||||
finish_reason="error",
|
||||
),
|
||||
LLMResponse(content="some other error", finish_reason="error"),
|
||||
])
|
||||
|
||||
response = await provider.chat_with_retry(messages=_IMAGE_MSG)
|
||||
|
||||
assert provider.calls == 2
|
||||
assert response.content == "some other error"
|
||||
assert response.finish_reason == "error"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_image_error_does_not_trigger_image_fallback() -> None:
|
||||
"""Regular non-transient errors must not trigger image stripping."""
|
||||
provider = ScriptedProvider([
|
||||
LLMResponse(content="401 unauthorized", finish_reason="error"),
|
||||
])
|
||||
|
||||
response = await provider.chat_with_retry(messages=_IMAGE_MSG)
|
||||
|
||||
assert provider.calls == 1
|
||||
assert response.content == "401 unauthorized"
|
||||
|
||||
Reference in New Issue
Block a user