178 lines
7.6 KiB
Python
178 lines
7.6 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
import pytest
|
|
|
|
from beaver.foundation.config.schema import AuthzConfig, BackendIdentityConfig, BeaverConfig
|
|
from beaver.tools.base import ObjectBackedTool, ToolContext
|
|
from beaver.tools.builtins import (
|
|
UserFilesCopyToWorkspaceTool,
|
|
UserFilesListTool,
|
|
UserFilesPublishOutputTool,
|
|
UserFilesReadTool,
|
|
UserFilesWriteTool,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_user_file_tools_write_read_and_list(tmp_path) -> None:
|
|
context = ToolContext(workspace=str(tmp_path))
|
|
write = ObjectBackedTool(UserFilesWriteTool())
|
|
read = ObjectBackedTool(UserFilesReadTool())
|
|
list_files = ObjectBackedTool(UserFilesListTool())
|
|
|
|
written = await write.invoke(
|
|
{"path": "outputs/summary.md", "content": "# Summary", "content_type": "text/markdown"},
|
|
context,
|
|
)
|
|
listed = await list_files.invoke({"path": "outputs"}, context)
|
|
loaded = await read.invoke({"path": "outputs/summary.md"}, context)
|
|
|
|
assert written.success is True
|
|
assert json.loads(written.content)["path"] == "outputs/summary.md"
|
|
assert listed.success is True
|
|
assert [item["name"] for item in json.loads(listed.content)["items"]] == ["summary.md"]
|
|
assert loaded.success is True
|
|
assert json.loads(loaded.content)["content"] == "# Summary"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_user_file_tools_reject_agent_write_to_uploads(tmp_path) -> None:
|
|
context = ToolContext(workspace=str(tmp_path))
|
|
write = ObjectBackedTool(UserFilesWriteTool())
|
|
|
|
result = await write.invoke({"path": "uploads/notes.txt", "content": "changed"}, context)
|
|
|
|
assert result.success is False
|
|
assert "uploads/ is user-provided input storage" in (result.error or "")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_user_file_tools_enforce_current_task_namespace(tmp_path) -> None:
|
|
context = ToolContext(workspace=str(tmp_path), services={"task_id": "task-123"})
|
|
write = ObjectBackedTool(UserFilesWriteTool())
|
|
|
|
current = await write.invoke({"path": "tasks/task-123/drafts/notes.md", "content": "ok"}, context)
|
|
direct = await write.invoke({"path": "tasks/notes.md", "content": "bad"}, context)
|
|
other = await write.invoke({"path": "tasks/task-456/notes.md", "content": "bad"}, context)
|
|
|
|
assert current.success is True
|
|
assert direct.success is False
|
|
assert "tasks/task-123/" in (direct.error or "")
|
|
assert other.success is False
|
|
assert "tasks/task-123/" in (other.error or "")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_user_file_tools_allow_shared_context_write(tmp_path) -> None:
|
|
context = ToolContext(workspace=str(tmp_path), services={"task_id": "task-123"})
|
|
write = ObjectBackedTool(UserFilesWriteTool())
|
|
read = ObjectBackedTool(UserFilesReadTool())
|
|
|
|
written = await write.invoke({"path": "shared/profile.json", "content": "{\"name\":\"Alice\"}"}, context)
|
|
loaded = await read.invoke({"path": "shared/profile.json"}, context)
|
|
|
|
assert written.success is True
|
|
assert loaded.success is True
|
|
assert json.loads(loaded.content)["content"] == "{\"name\":\"Alice\"}"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_user_file_tools_copy_to_workspace_and_publish_output(tmp_path) -> None:
|
|
uploads_dir = tmp_path / "user_files" / "uploads"
|
|
uploads_dir.mkdir(parents=True)
|
|
(uploads_dir / "get_helm.sh").write_text(": ${USE_SUDO:=\"true\"}\n", encoding="utf-8")
|
|
context = ToolContext(
|
|
workspace=str(tmp_path),
|
|
services={"task_id": "task-123"},
|
|
metadata={"run_id": "run-1"},
|
|
)
|
|
copy_tool = ObjectBackedTool(UserFilesCopyToWorkspaceTool())
|
|
publish_tool = ObjectBackedTool(UserFilesPublishOutputTool())
|
|
read = ObjectBackedTool(UserFilesReadTool())
|
|
|
|
copied = await copy_tool.invoke({"path": "uploads/get_helm.sh"}, context)
|
|
copied_payload = json.loads(copied.content)
|
|
staged = tmp_path / copied_payload["workspace_path"]
|
|
staged.write_text(": ${USE_SUDO:=\"false\"}\n", encoding="utf-8")
|
|
published = await publish_tool.invoke(
|
|
{"source_path": copied_payload["workspace_path"], "target_path": "outputs/get_helm.no-sudo.sh"},
|
|
context,
|
|
)
|
|
original = await read.invoke({"path": "uploads/get_helm.sh"}, context)
|
|
output = await read.invoke({"path": "outputs/get_helm.no-sudo.sh"}, context)
|
|
|
|
assert copied.success is True
|
|
assert copied_payload["workspace_path"] == "user-files/tasks/task-123/get_helm.sh"
|
|
assert published.success is True
|
|
assert json.loads(original.content)["content"] == ": ${USE_SUDO:=\"true\"}\n"
|
|
assert json.loads(output.content)["content"] == ": ${USE_SUDO:=\"false\"}\n"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_user_file_publish_rejects_non_output_target_and_workspace_escape(tmp_path) -> None:
|
|
context = ToolContext(workspace=str(tmp_path))
|
|
source = tmp_path / "result.txt"
|
|
source.write_text("done", encoding="utf-8")
|
|
outside = tmp_path.parent / "outside.txt"
|
|
outside.write_text("outside", encoding="utf-8")
|
|
publish_tool = ObjectBackedTool(UserFilesPublishOutputTool())
|
|
|
|
upload_target = await publish_tool.invoke({"source_path": "result.txt", "target_path": "uploads/result.txt"}, context)
|
|
escaped_source = await publish_tool.invoke({"source_path": str(outside), "target_path": "outputs/result.txt"}, context)
|
|
|
|
assert upload_target.success is False
|
|
assert "outputs/" in (upload_target.error or "")
|
|
assert escaped_source.success is False
|
|
assert "escapes workspace" in (escaped_source.error or "")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_user_file_tools_reject_internal_workspace_paths(tmp_path) -> None:
|
|
context = ToolContext(workspace=str(tmp_path))
|
|
read = ObjectBackedTool(UserFilesReadTool())
|
|
write = ObjectBackedTool(UserFilesWriteTool())
|
|
|
|
read_result = await read.invoke({"path": "uploads/../../state/secrets.json"}, context)
|
|
write_result = await write.invoke({"path": "workspace/debug.txt", "content": "x"}, context)
|
|
|
|
assert read_result.success is False
|
|
assert "Parent-directory traversal" in read_result.error
|
|
assert write_result.success is False
|
|
assert "Path must be under" in write_result.error
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_user_file_tools_reject_absolute_style_user_paths(tmp_path) -> None:
|
|
context = ToolContext(workspace=str(tmp_path), services={"task_id": "task-123"})
|
|
read = ObjectBackedTool(UserFilesReadTool())
|
|
write = ObjectBackedTool(UserFilesWriteTool())
|
|
list_files = ObjectBackedTool(UserFilesListTool())
|
|
|
|
read_result = await read.invoke({"path": "/uploads/input.txt"}, context)
|
|
write_result = await write.invoke({"path": "/outputs/result.txt", "content": "x"}, context)
|
|
task_write = await write.invoke({"path": "/tasks/task-123/draft.md", "content": "x"}, context)
|
|
list_result = await list_files.invoke({"path": "/shared/profile.json"}, context)
|
|
|
|
for result in (read_result, write_result, task_write, list_result):
|
|
assert result.success is False
|
|
assert "Absolute paths are not allowed" in (result.error or "")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_user_file_tools_report_missing_deployed_minio_settings(tmp_path, monkeypatch) -> None:
|
|
monkeypatch.delenv("BEAVER_AUTHZ_INTERNAL_TOKEN", raising=False)
|
|
monkeypatch.delenv("AUTHZ_INTERNAL_TOKEN", raising=False)
|
|
config = BeaverConfig(
|
|
authz=AuthzConfig(enabled=True, base_url="http://authz.local"),
|
|
backend_identity=BackendIdentityConfig(backend_id="alice", client_id="alice", client_secret="secret"),
|
|
)
|
|
context = ToolContext(workspace=str(tmp_path), services={"beaver_config": config})
|
|
write = ObjectBackedTool(UserFilesWriteTool())
|
|
|
|
result = await write.invoke({"path": "outputs/summary.md", "content": "# Summary"}, context)
|
|
|
|
assert result.success is False
|
|
assert "AuthZ internal token is not configured" in (result.error or "")
|