diff --git a/backend/app/services/agent_tools.py b/backend/app/services/agent_tools.py index 392a9b557..328f33cd6 100644 --- a/backend/app/services/agent_tools.py +++ b/backend/app/services/agent_tools.py @@ -123,8 +123,8 @@ _settings = get_settings() WORKSPACE_ROOT = Path(_settings.STORAGE_LOCAL_ROOT or _settings.AGENT_DATA_DIR) -TOOL_MATERIALIZE_MAX_FILE_BYTES = 10 * 1024 * 1024 -TOOL_MATERIALIZE_MAX_TOTAL_BYTES = 100 * 1024 * 1024 +TOOL_MATERIALIZE_MAX_FILE_BYTES = 50 * 1024 * 1024 +TOOL_MATERIALIZE_MAX_TOTAL_BYTES = 500 * 1024 * 1024 FEISHU_APPROVAL_ATTACHMENT_MAX_BYTES = 50 * 1024 * 1024 FEISHU_APPROVAL_IMAGE_MAX_BYTES = 10 * 1024 * 1024 FEISHU_APPROVAL_CODE_MAX_CHARS = 256 @@ -1661,8 +1661,24 @@ async def _materialize_storage_path_with_budget( if await storage.is_file(storage_key): version = await storage.get_version(storage_key) if version.size > max_file_bytes: + logger.warning( + "Tool workspace materialization skipped file: " + "path={} size_bytes={} limit_bytes={} reason={}", + rel_path, + version.size, + max_file_bytes, + "per_file_limit", + ) return if budget["total"] + version.size > TOOL_MATERIALIZE_MAX_TOTAL_BYTES: + logger.warning( + "Tool workspace materialization skipped file: " + "path={} size_bytes={} limit_bytes={} reason={}", + rel_path, + version.size, + TOOL_MATERIALIZE_MAX_TOTAL_BYTES, + "total_limit", + ) return target = (local_root / rel_path).resolve() if not target.is_relative_to(local_root.resolve()): diff --git a/backend/tests/test_agent_tools_storage_workspace.py b/backend/tests/test_agent_tools_storage_workspace.py index d660be81b..46ba5d9d5 100644 --- a/backend/tests/test_agent_tools_storage_workspace.py +++ b/backend/tests/test_agent_tools_storage_workspace.py @@ -1,4 +1,5 @@ from contextlib import asynccontextmanager +from unittest.mock import AsyncMock, Mock import uuid import pytest @@ -165,6 +166,94 @@ async def test_temp_workspace_materializes_only_requested_paths(monkeypatch): temp_ws.cleanup() +def test_temp_workspace_materialization_limits_are_50_and_500_mib(): + assert agent_tools.TOOL_MATERIALIZE_MAX_FILE_BYTES == 50 * 1024 * 1024 + assert agent_tools.TOOL_MATERIALIZE_MAX_TOTAL_BYTES == 500 * 1024 * 1024 + + +@pytest.mark.asyncio +async def test_temp_workspace_materializes_file_above_previous_10_mib_limit( + monkeypatch, +): + agent_id = uuid.uuid4() + content = b"x" * (11 * 1024 * 1024) + storage = MemoryStorageBackend({ + f"{agent_id}/workspace/presentation.pptx": content, + }) + monkeypatch.setattr(agent_tools, "get_storage_backend", lambda: storage) + + temp_ws = await agent_tools._prepare_temp_workspace( + agent_id, + paths=["workspace/presentation.pptx"], + ) + try: + assert (temp_ws.root / "workspace" / "presentation.pptx").read_bytes() == content + finally: + temp_ws.cleanup() + + +@pytest.mark.asyncio +async def test_temp_workspace_logs_file_skipped_by_per_file_limit(monkeypatch, tmp_path): + agent_id = uuid.uuid4() + storage_key = f"{agent_id}/workspace/oversized.pptx" + storage = MemoryStorageBackend({storage_key: b"too large"}) + storage.get_version = AsyncMock( # type: ignore[method-assign] + return_value=StorageVersion( + key=storage_key, + exists=True, + is_dir=False, + size=51 * 1024 * 1024, + ) + ) + warning = Mock() + monkeypatch.setattr(agent_tools.logger, "warning", warning) + + await agent_tools._materialize_storage_path_with_budget( + storage, + storage_key, + "workspace/oversized.pptx", + tmp_path, + {"total": 0}, + {}, + ) + + assert not (tmp_path / "workspace" / "oversized.pptx").exists() + warning.assert_called_once_with( + "Tool workspace materialization skipped file: path={} size_bytes={} limit_bytes={} reason={}", + "workspace/oversized.pptx", + 51 * 1024 * 1024, + 50 * 1024 * 1024, + "per_file_limit", + ) + + +@pytest.mark.asyncio +async def test_temp_workspace_logs_file_skipped_by_total_limit(monkeypatch, tmp_path): + agent_id = uuid.uuid4() + storage_key = f"{agent_id}/workspace/second.pptx" + storage = MemoryStorageBackend({storage_key: b"second"}) + warning = Mock() + monkeypatch.setattr(agent_tools.logger, "warning", warning) + + await agent_tools._materialize_storage_path_with_budget( + storage, + storage_key, + "workspace/second.pptx", + tmp_path, + {"total": agent_tools.TOOL_MATERIALIZE_MAX_TOTAL_BYTES}, + {}, + ) + + assert not (tmp_path / "workspace" / "second.pptx").exists() + warning.assert_called_once_with( + "Tool workspace materialization skipped file: path={} size_bytes={} limit_bytes={} reason={}", + "workspace/second.pptx", + len(b"second"), + 500 * 1024 * 1024, + "total_limit", + ) + + @pytest.mark.asyncio async def test_execute_tool_list_files_does_not_create_persistent_workspace(monkeypatch, tmp_path): agent_id = uuid.uuid4()