mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-09 04:40:57 +00:00
fix(runtime): bound tenant resource amplification
This commit is contained in:
@@ -50,11 +50,12 @@ def _create_mock_webhook(
|
||||
return webhook
|
||||
|
||||
|
||||
def _create_mock_result(items: list = None, first_item=None):
|
||||
def _create_mock_result(items: list = None, first_item=None, scalar_value=None):
|
||||
"""Create mock result object for persistence queries."""
|
||||
result = Mock()
|
||||
result.all = Mock(return_value=items or [])
|
||||
result.first = Mock(return_value=first_item)
|
||||
result.scalar = Mock(return_value=scalar_value)
|
||||
result.rowcount = 1
|
||||
return result
|
||||
|
||||
@@ -154,6 +155,8 @@ class TestWebhookServiceCreateWebhook:
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
return _create_mock_result(scalar_value=0) # Count
|
||||
if call_count == 2:
|
||||
return insert_result # Insert
|
||||
return select_result # Select
|
||||
|
||||
@@ -205,6 +208,8 @@ class TestWebhookServiceCreateWebhook:
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
return _create_mock_result(scalar_value=0)
|
||||
if call_count == 2:
|
||||
return _create_write_result() # Insert
|
||||
return _create_mock_result(first_item=created_webhook)
|
||||
|
||||
@@ -246,6 +251,8 @@ class TestWebhookServiceCreateWebhook:
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
return _create_mock_result(scalar_value=0)
|
||||
if call_count == 2:
|
||||
return _create_write_result()
|
||||
return _create_mock_result(first_item=created_webhook)
|
||||
|
||||
@@ -265,6 +272,42 @@ class TestWebhookServiceCreateWebhook:
|
||||
# Verify
|
||||
assert result['enabled'] is False
|
||||
|
||||
async def test_create_webhook_rejects_workspace_at_capacity(self):
|
||||
ap = SimpleNamespace(
|
||||
instance_config=SimpleNamespace(
|
||||
data={'webhooks': {'max_per_workspace': 2}},
|
||||
),
|
||||
persistence_mgr=SimpleNamespace(
|
||||
execute_async=AsyncMock(return_value=_create_mock_result(scalar_value=2)),
|
||||
),
|
||||
)
|
||||
|
||||
service = WebhookService(ap)
|
||||
|
||||
with pytest.raises(ValueError, match=r'Maximum number of webhooks \(2\) reached'):
|
||||
await service.create_webhook(
|
||||
WORKSPACE_UUID,
|
||||
name='Too many',
|
||||
url='https://example.invalid',
|
||||
)
|
||||
|
||||
ap.persistence_mgr.execute_async.assert_awaited_once()
|
||||
|
||||
async def test_max_per_workspace_clamps_invalid_and_oversized_values(self):
|
||||
ap = SimpleNamespace(
|
||||
instance_config=SimpleNamespace(
|
||||
data={'webhooks': {'max_per_workspace': 999999}},
|
||||
)
|
||||
)
|
||||
service = WebhookService(ap)
|
||||
assert service.max_per_workspace() == 64
|
||||
|
||||
ap.instance_config.data['webhooks']['max_per_workspace'] = 0
|
||||
assert service.max_per_workspace() == 1
|
||||
|
||||
ap.instance_config.data['webhooks']['max_per_workspace'] = 'invalid'
|
||||
assert service.max_per_workspace() == 16
|
||||
|
||||
|
||||
class TestWebhookServiceGetWebhook:
|
||||
"""Tests for get_webhook method."""
|
||||
|
||||
@@ -1280,6 +1280,45 @@ async def test_box_service_rejects_and_cleans_up_when_execution_exceeds_workspac
|
||||
assert backend.stop_calls == ['person_test_user']
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_box_service_rejects_workspace_inode_bomb_before_execution(tmp_path):
|
||||
logger = Mock()
|
||||
backend = FakeBackend(logger)
|
||||
runtime = BoxRuntime(logger=logger, backends=[backend], session_ttl_sec=300)
|
||||
host_dir = tmp_path / 'quota-workspace-entries'
|
||||
host_dir.mkdir()
|
||||
app = make_app(logger, [str(tmp_path)], workspace_quota_mb=1)
|
||||
app.instance_config.data['box']['local']['default_workspace'] = str(host_dir)
|
||||
app.instance_config.data['box']['limits'] = {'max_workspace_entries': 2}
|
||||
service = BoxService(app, client=_InProcessBoxRuntimeClient(logger, runtime))
|
||||
|
||||
tenant_host_dir = service._tenant_workspace(_CONTEXT)
|
||||
assert tenant_host_dir is not None
|
||||
os.makedirs(tenant_host_dir, exist_ok=True)
|
||||
for index in range(3):
|
||||
pathlib.Path(tenant_host_dir, f'tiny-{index}').write_bytes(b'x')
|
||||
|
||||
await service.initialize()
|
||||
|
||||
with pytest.raises(BoxValidationError, match='workspace entry limit exceeded before execution'):
|
||||
await service.execute_tool({'command': 'echo hi'}, make_query(46))
|
||||
|
||||
assert backend.start_calls == []
|
||||
|
||||
|
||||
def test_box_service_workspace_entry_limit_is_hard_clamped():
|
||||
app = make_app(Mock())
|
||||
app.instance_config.data['box']['limits'] = {'max_workspace_entries': 10_000_000}
|
||||
service = BoxService(app, client=Mock(spec=BoxRuntimeClient))
|
||||
assert service._max_workspace_entries() == 1_000_000
|
||||
|
||||
app.instance_config.data['box']['limits']['max_workspace_entries'] = 0
|
||||
assert service._max_workspace_entries() == 1
|
||||
|
||||
app.instance_config.data['box']['limits']['max_workspace_entries'] = 'invalid'
|
||||
assert service._max_workspace_entries() == 100_000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_profile_offline_readonly_locks_read_only_rootfs():
|
||||
"""offline_readonly locks read_only_rootfs so it cannot be overridden."""
|
||||
@@ -2126,7 +2165,7 @@ class TestInboundOutboundRoundTrip:
|
||||
|
||||
async def fake_execute_tool(parameters, q):
|
||||
calls.append(parameters['command'])
|
||||
if 'os.walk' in parameters['command']:
|
||||
if 'os.scandir' in parameters['command']:
|
||||
return {
|
||||
'ok': True,
|
||||
'stdout': '[{"name": "out.png", "b64": "QUJD"}]',
|
||||
@@ -2156,7 +2195,7 @@ class TestInboundOutboundRoundTrip:
|
||||
|
||||
async def fake_execute_tool(parameters, q):
|
||||
calls.append(parameters['command'])
|
||||
if 'os.walk' in parameters['command']:
|
||||
if 'os.scandir' in parameters['command']:
|
||||
return {'ok': True, 'stdout': '[]', 'stderr': ''}
|
||||
return {'ok': True, 'stdout': '', 'stderr': ''}
|
||||
|
||||
|
||||
@@ -65,6 +65,9 @@ def test_wrap_python_command_with_env_contains_bootstrap_and_command():
|
||||
assert '_LB_SYSTEM_PYTHON="$(command -v python3 || command -v python || true)"' in command
|
||||
assert '"$_LB_SYSTEM_PYTHON" -m venv "$_LB_VENV_DIR"' in command
|
||||
assert 'kill -0 "$_LB_LOCK_OWNER"' in command
|
||||
assert 'max_manifest_bytes = 10 * 1024 * 1024' in command
|
||||
assert 'handle.read(1024 * 1024)' in command
|
||||
assert 'digest.update(handle.read())' not in command
|
||||
assert 'export VIRTUAL_ENV="$_LB_VENV_DIR"' in command
|
||||
assert command.rstrip().endswith('python script.py')
|
||||
|
||||
|
||||
@@ -1,6 +1,11 @@
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
import langbot_plugin.api.entities.builtin.platform.message as platform_message
|
||||
from langbot.pkg.pipeline.longtext.strategies.image import Text2ImageStrategy
|
||||
from langbot.pkg.pipeline.longtext.strategies import image
|
||||
|
||||
|
||||
class _WideFont:
|
||||
@@ -15,3 +20,31 @@ def test_image_strategy_line_split_always_consumes_input():
|
||||
|
||||
assert lines == ['a', 'b', 'c']
|
||||
assert ''.join(lines) == 'abc'
|
||||
|
||||
|
||||
def test_image_strategy_numeric_boundaries_are_found_in_linear_order():
|
||||
strategy = Text2ImageStrategy(Mock())
|
||||
|
||||
assert strategy.indexNumber('a12-b12-c345') == [['12', 1], ['12', 5], ['345', 9]]
|
||||
|
||||
|
||||
def test_image_strategy_rejects_unbounded_line_count_before_allocating_canvas(monkeypatch):
|
||||
strategy = Text2ImageStrategy(Mock())
|
||||
monkeypatch.setattr(image, '_MAX_TEXT_TO_IMAGE_LINES', 2)
|
||||
|
||||
with pytest.raises(ValueError, match='2 lines'):
|
||||
strategy._split_text_lines('one\ntwo\nthree', 1000, _WideFont())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_strategy_falls_back_to_forward_for_oversized_text(monkeypatch):
|
||||
app = Mock()
|
||||
strategy = Text2ImageStrategy(app)
|
||||
monkeypatch.setattr(image, '_MAX_TEXT_TO_IMAGE_CHARS', 4)
|
||||
query = SimpleNamespace(adapter=SimpleNamespace(bot_account_id='bot'))
|
||||
|
||||
components = await strategy.process('12345', query)
|
||||
|
||||
assert len(components) == 1
|
||||
assert isinstance(components[0], platform_message.Forward)
|
||||
app.logger.warning.assert_called_once()
|
||||
|
||||
@@ -0,0 +1,99 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.platform.webhook_pusher import WebhookPusher
|
||||
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
def _application(max_inflight_requests: object) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
instance_config=SimpleNamespace(
|
||||
data={
|
||||
'webhooks': {
|
||||
'max_inflight_requests': max_inflight_requests,
|
||||
}
|
||||
}
|
||||
),
|
||||
logger=logging.getLogger(__name__),
|
||||
)
|
||||
|
||||
|
||||
async def test_delivery_admission_never_queues_above_instance_limit():
|
||||
pusher = WebhookPusher(_application(2))
|
||||
release = asyncio.Event()
|
||||
both_started = asyncio.Event()
|
||||
calls = 0
|
||||
active = 0
|
||||
peak_active = 0
|
||||
|
||||
async def fake_push(url: str, payload: dict) -> dict:
|
||||
nonlocal calls, active, peak_active
|
||||
calls += 1
|
||||
active += 1
|
||||
peak_active = max(peak_active, active)
|
||||
if active == 2:
|
||||
both_started.set()
|
||||
try:
|
||||
await release.wait()
|
||||
return {'url': url}
|
||||
finally:
|
||||
active -= 1
|
||||
|
||||
pusher._push_to_webhook = fake_push
|
||||
webhooks = [{'url': f'https://example.invalid/{index}'} for index in range(5)]
|
||||
|
||||
first_delivery = asyncio.create_task(pusher._push_to_webhooks(webhooks, {}))
|
||||
await asyncio.wait_for(both_started.wait(), timeout=1)
|
||||
second_results = await pusher._push_to_webhooks(webhooks, {})
|
||||
release.set()
|
||||
first_results = await first_delivery
|
||||
|
||||
assert len(first_results) == 2
|
||||
assert second_results == []
|
||||
assert calls == 2
|
||||
assert peak_active == 2
|
||||
assert pusher._inflight_requests == 0
|
||||
|
||||
|
||||
async def test_cancelled_delivery_reaps_children_and_releases_slots():
|
||||
pusher = WebhookPusher(_application(1))
|
||||
started = asyncio.Event()
|
||||
never = asyncio.Event()
|
||||
|
||||
async def blocking_push(url: str, payload: dict) -> dict:
|
||||
started.set()
|
||||
await never.wait()
|
||||
return {}
|
||||
|
||||
pusher._push_to_webhook = blocking_push
|
||||
delivery = asyncio.create_task(
|
||||
pusher._push_to_webhooks([{'url': 'https://example.invalid'}], {}),
|
||||
)
|
||||
await asyncio.wait_for(started.wait(), timeout=1)
|
||||
|
||||
delivery.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await delivery
|
||||
|
||||
assert pusher._inflight_requests == 0
|
||||
pusher._push_to_webhook = AsyncMock(return_value={})
|
||||
assert await pusher._push_to_webhooks([{'url': 'https://example.invalid'}], {}) == [{}]
|
||||
|
||||
|
||||
async def test_max_inflight_requests_clamps_config():
|
||||
pusher = WebhookPusher(_application(999999))
|
||||
assert pusher._max_inflight_requests() == 128
|
||||
|
||||
pusher.ap.instance_config.data['webhooks']['max_inflight_requests'] = 0
|
||||
assert pusher._max_inflight_requests() == 1
|
||||
|
||||
pusher.ap.instance_config.data['webhooks']['max_inflight_requests'] = 'invalid'
|
||||
assert pusher._max_inflight_requests() == 16
|
||||
@@ -37,6 +37,7 @@ def _make_adapter(load_return=b'hello', load_side_effect=None):
|
||||
provider.load = AsyncMock(return_value=load_return, side_effect=load_side_effect)
|
||||
storage_mgr = Mock()
|
||||
storage_mgr.storage_provider = provider
|
||||
storage_mgr.load_scoped_object_key = AsyncMock(return_value=load_return, side_effect=load_side_effect)
|
||||
storage_mgr.scoped_prefix.return_value = _UPLOAD_PREFIX
|
||||
storage_mgr.is_scoped_object_key.return_value = True
|
||||
storage_mgr.delete_scoped_object_key = AsyncMock()
|
||||
@@ -94,7 +95,7 @@ async def test_file_uses_octet_stream_fallback():
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skips_components_without_path_or_unknown_type():
|
||||
adapter, _, provider = _make_adapter()
|
||||
adapter, storage_mgr, provider = _make_adapter()
|
||||
chain = [
|
||||
{'type': 'Image', 'path': ''}, # no path
|
||||
{'type': 'Plain', 'path': 'storage://abc/x'}, # not a file component
|
||||
@@ -102,6 +103,7 @@ async def test_skips_components_without_path_or_unknown_type():
|
||||
]
|
||||
await adapter._process_image_components(_make_connection(), chain)
|
||||
provider.load.assert_not_awaited()
|
||||
storage_mgr.load_scoped_object_key.assert_not_awaited()
|
||||
assert 'base64' not in chain[0]
|
||||
assert 'base64' not in chain[1]
|
||||
|
||||
|
||||
@@ -327,7 +327,7 @@ async def test_attachment_key_must_belong_to_connection_upload_scope():
|
||||
storage_mgr = Mock()
|
||||
storage_mgr.scoped_prefix.return_value = 'v1/current/upload_image/'
|
||||
storage_mgr.is_scoped_object_key.return_value = True
|
||||
storage_mgr.storage_provider.load = AsyncMock(return_value=b'image')
|
||||
storage_mgr.load_scoped_object_key = AsyncMock(return_value=b'image')
|
||||
storage_mgr.delete_scoped_object_key = AsyncMock()
|
||||
adapter = WebSocketAdapter.model_construct(
|
||||
ap=Mock(storage_mgr=storage_mgr),
|
||||
@@ -347,6 +347,11 @@ async def test_attachment_key_must_belong_to_connection_upload_scope():
|
||||
'v1/current/upload_image/key.png',
|
||||
expected_owner_type='upload_image',
|
||||
)
|
||||
storage_mgr.load_scoped_object_key.assert_awaited_once_with(
|
||||
connection.execution_context,
|
||||
'v1/current/upload_image/key.png',
|
||||
expected_owner_type='upload_image',
|
||||
)
|
||||
storage_mgr.delete_scoped_object_key.assert_awaited_once_with(
|
||||
connection.execution_context,
|
||||
'v1/current/upload_image/key.png',
|
||||
|
||||
@@ -417,7 +417,9 @@ class TestGetConfigFile:
|
||||
mock_app.persistence_mgr = Mock()
|
||||
mock_app.persistence_mgr.execute_async = AsyncMock()
|
||||
mock_app.storage_mgr = StorageMgr(mock_app)
|
||||
mock_app.storage_mgr.storage_provider = SimpleNamespace(load=AsyncMock(return_value=b'plugin config bytes'))
|
||||
mock_app.storage_mgr.storage_provider = SimpleNamespace(
|
||||
load_bounded=AsyncMock(return_value=b'plugin config bytes')
|
||||
)
|
||||
mock_app.logger = Mock()
|
||||
return mock_app
|
||||
|
||||
@@ -461,7 +463,10 @@ class TestGetConfigFile:
|
||||
|
||||
assert response.code == 0
|
||||
assert base64.b64decode(response.data['file_base64']) == b'plugin config bytes'
|
||||
app.storage_mgr.storage_provider.load.assert_awaited_once_with(file_key)
|
||||
app.storage_mgr.storage_provider.load_bounded.assert_awaited_once_with(
|
||||
file_key,
|
||||
max_bytes=10 * 1024 * 1024,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
@@ -478,7 +483,7 @@ class TestGetConfigFile:
|
||||
|
||||
assert response.code != 0
|
||||
assert 'Failed to load config file' in response.message
|
||||
app.storage_mgr.storage_provider.load.assert_not_awaited()
|
||||
app.storage_mgr.storage_provider.load_bounded.assert_not_awaited()
|
||||
|
||||
|
||||
class TestGetBinaryStorage:
|
||||
|
||||
@@ -584,6 +584,40 @@ async def test_host_file_api_falls_back_to_tenant_box_when_openat_is_unavailable
|
||||
assert 'path = "/workspace/file.txt"' in command
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_box_workspace_edit_script_bounds_file_read_and_replacement(monkeypatch):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
box_service = SimpleNamespace(
|
||||
available=True,
|
||||
default_workspace=tmpdir,
|
||||
_tenant_workspace=Mock(return_value=tmpdir),
|
||||
execute_tool=AsyncMock(
|
||||
return_value={
|
||||
'ok': True,
|
||||
'stdout': '{"ok": false, "error": "File exceeds limit"}',
|
||||
'stderr': '',
|
||||
}
|
||||
),
|
||||
)
|
||||
loader = NativeToolLoader(SimpleNamespace(box_service=box_service, logger=Mock()))
|
||||
monkeypatch.setattr(native_loader, '_SECURE_HOST_FILE_OPS_AVAILABLE', False)
|
||||
|
||||
await loader.invoke_tool(
|
||||
'edit',
|
||||
{
|
||||
'path': '/workspace/file.txt',
|
||||
'old_string': 'old',
|
||||
'new_string': 'new',
|
||||
},
|
||||
_make_query(),
|
||||
)
|
||||
|
||||
command = box_service.execute_tool.await_args.args[0]['command']
|
||||
assert f'os.path.getsize(path) > {native_loader._MAX_HOST_EDIT_FILE_BYTES}' in command
|
||||
assert f'f.read({native_loader._MAX_HOST_EDIT_FILE_BYTES + 1})' in command
|
||||
assert f"len(new_content.encode('utf-8')) > {native_loader._MAX_HOST_EDIT_FILE_BYTES}" in command
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_box_availability_helper_handles_unavailable_and_errors():
|
||||
from langbot.pkg.provider.tools.loaders.availability import is_box_backend_available
|
||||
|
||||
@@ -52,6 +52,7 @@ def _make_app() -> Mock:
|
||||
storage_mgr.storage_provider = Mock()
|
||||
storage_mgr.storage_provider.exists = AsyncMock(return_value=True)
|
||||
storage_mgr.storage_provider.load = AsyncMock()
|
||||
storage_mgr.storage_provider.load_bounded = AsyncMock()
|
||||
storage_mgr.storage_provider.save = AsyncMock()
|
||||
storage_mgr.storage_provider.size = AsyncMock(return_value=123)
|
||||
storage_mgr.storage_provider.delete = AsyncMock()
|
||||
@@ -175,7 +176,7 @@ class TestStoreZipFile:
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_zip_file_extracts_supported_files_and_skips_noise(self):
|
||||
kb = _make_kb()
|
||||
kb.ap.storage_mgr.storage_provider.load = AsyncMock(
|
||||
kb.ap.storage_mgr.storage_provider.load_bounded = AsyncMock(
|
||||
return_value=_make_zip_bytes(
|
||||
{
|
||||
'doc1.pdf': b'pdf',
|
||||
@@ -210,7 +211,7 @@ class TestStoreZipFile:
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_zip_file_raises_when_no_supported_files(self):
|
||||
kb = _make_kb()
|
||||
kb.ap.storage_mgr.storage_provider.load = AsyncMock(
|
||||
kb.ap.storage_mgr.storage_provider.load_bounded = AsyncMock(
|
||||
return_value=_make_zip_bytes({'image.png': b'png', 'video.mp4': b'video'})
|
||||
)
|
||||
kb.store_file = AsyncMock()
|
||||
@@ -224,7 +225,7 @@ class TestStoreZipFile:
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_zip_file_rejects_too_many_documents_before_extracting(self):
|
||||
kb = _make_kb()
|
||||
kb.ap.storage_mgr.storage_provider.load = AsyncMock(
|
||||
kb.ap.storage_mgr.storage_provider.load_bounded = AsyncMock(
|
||||
return_value=_make_zip_bytes({f'doc-{index}.txt': b'text' for index in range(9)})
|
||||
)
|
||||
kb.store_file = AsyncMock()
|
||||
@@ -239,7 +240,7 @@ class TestStoreZipFile:
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_zip_file_rejects_extreme_compression_ratio_before_extracting(self):
|
||||
kb = _make_kb()
|
||||
kb.ap.storage_mgr.storage_provider.load = AsyncMock(
|
||||
kb.ap.storage_mgr.storage_provider.load_bounded = AsyncMock(
|
||||
return_value=_make_zip_bytes(
|
||||
{'bomb.txt': b'A' * (1024 * 1024)},
|
||||
compression=zipfile.ZIP_DEFLATED,
|
||||
|
||||
@@ -168,6 +168,15 @@ class TestPathTraversalPrevention:
|
||||
assert loaded == content
|
||||
await provider.delete(key)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bounded_load_stops_after_limit(self, storage_provider):
|
||||
provider, storage_path = storage_provider
|
||||
|
||||
with patch('langbot.pkg.storage.providers.localstorage.LOCAL_STORAGE_PATH', storage_path):
|
||||
await provider.save('oversized.bin', b'12345')
|
||||
with pytest.raises(ValueError, match='4-byte read limit'):
|
||||
await provider.load_bounded('oversized.bin', max_bytes=4)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_dir_recursive_non_existing_dir(self, storage_provider):
|
||||
"""delete_dir_recursive should handle non-existing directories gracefully."""
|
||||
|
||||
@@ -134,6 +134,16 @@ class TestS3StorageProviderWithMoto:
|
||||
loaded_data = await provider.load('test/file.txt')
|
||||
assert loaded_data == test_data
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bounded_load_rejects_oversized_object(self, mock_app_with_s3_config, s3_mock):
|
||||
s3storage = get_s3storage_module()
|
||||
provider = s3storage.S3StorageProvider(mock_app_with_s3_config)
|
||||
await provider.initialize()
|
||||
await provider.save('test/oversized.bin', b'12345')
|
||||
|
||||
with pytest.raises(ValueError, match='4-byte read limit'):
|
||||
await provider.load_bounded('test/oversized.bin', max_bytes=4)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_exists_returns_true_for_existing_object(self, mock_app_with_s3_config, s3_mock):
|
||||
"""Test that exists returns True for existing object."""
|
||||
|
||||
@@ -8,6 +8,7 @@ import pytest
|
||||
from langbot.pkg.api.http.authz import WorkspaceRequiredError
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
from langbot.pkg.storage.mgr import StorageMgr
|
||||
from langbot.pkg.storage.provider import HARD_MAX_STORAGE_OBJECT_BYTES
|
||||
from langbot.pkg.utils.bounded_executor import current_blocking_work_scope
|
||||
|
||||
|
||||
@@ -202,6 +203,39 @@ async def test_opaque_object_operations_reject_cross_scope_and_owner_type(manage
|
||||
assert object_key in manager.storage_provider.values
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scoped_object_reads_and_writes_have_configured_and_hard_byte_limits(manager):
|
||||
manager.ap.instance_config = SimpleNamespace(data={'storage': {'max_object_read_bytes': 4}})
|
||||
|
||||
with pytest.raises(ValueError, match='4-byte write limit'):
|
||||
await manager.save_scoped(
|
||||
_context(WORKSPACE_A),
|
||||
owner_type='upload',
|
||||
owner='account:a',
|
||||
key='oversized.bin',
|
||||
value=b'12345',
|
||||
)
|
||||
|
||||
object_key = manager.scoped_object_key(
|
||||
_context(WORKSPACE_A),
|
||||
owner_type='upload',
|
||||
owner='account:a',
|
||||
key='external.bin',
|
||||
)
|
||||
manager.storage_provider.values[object_key] = b'12345'
|
||||
manager.storage_provider.load = AsyncMock(side_effect=AssertionError('oversized object must not be loaded'))
|
||||
with pytest.raises(ValueError, match='4-byte read limit'):
|
||||
await manager.load_scoped_object_key(
|
||||
_context(WORKSPACE_A),
|
||||
object_key,
|
||||
expected_owner_type='upload',
|
||||
)
|
||||
manager.storage_provider.load.assert_not_awaited()
|
||||
|
||||
manager.ap.instance_config.data['storage']['max_object_read_bytes'] = HARD_MAX_STORAGE_OBJECT_BYTES + 1
|
||||
assert manager._object_read_limit() == HARD_MAX_STORAGE_OBJECT_BYTES
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scoped_provider_is_not_touched_after_generation_is_fenced(manager):
|
||||
object_key = await manager.save_scoped(
|
||||
|
||||
@@ -268,9 +268,9 @@ class TestDeleteByFilterGuard:
|
||||
backend.ap = type('Ap', (), {'logger': AsyncMock()})()
|
||||
backend._ensure_client = AsyncMock(return_value=backend._client)
|
||||
backend._index_exists = AsyncMock(return_value=True)
|
||||
# _search_keys must never be reached for an unusable filter.
|
||||
backend._search_keys = AsyncMock(
|
||||
side_effect=AssertionError('_search_keys must not be called for an unusable filter')
|
||||
# The deletion scan must never be reached for an unusable filter.
|
||||
backend._delete_search_results = AsyncMock(
|
||||
side_effect=AssertionError('_delete_search_results must not be called for an unusable filter')
|
||||
)
|
||||
|
||||
# Filter references only a non-indexed field -> maps to no FT conditions.
|
||||
@@ -285,12 +285,57 @@ class TestDeleteByFilterGuard:
|
||||
backend.ap = type('Ap', (), {'logger': AsyncMock()})()
|
||||
backend._ensure_client = AsyncMock(return_value=backend._client)
|
||||
backend._index_exists = AsyncMock(return_value=True)
|
||||
backend._search_keys = AsyncMock(return_value=['kb:col1:id1', 'kb:col1:id2'])
|
||||
backend._delete_search_results = AsyncMock(return_value=2)
|
||||
|
||||
deleted = await backend.delete_by_filter('col1', {'file_id': 'f1'})
|
||||
|
||||
assert deleted == 2
|
||||
backend._client.delete.assert_awaited_once_with(['kb:col1:id1', 'kb:col1:id2'])
|
||||
backend._delete_search_results.assert_awaited_once_with(
|
||||
backend._client,
|
||||
backend._index_name('col1'),
|
||||
'@file_id:{f1}',
|
||||
)
|
||||
|
||||
|
||||
class TestBatchedDelete:
|
||||
async def test_matching_keys_are_deleted_in_fixed_pages(self, monkeypatch):
|
||||
mod = get_valkey_module()
|
||||
backend = make_backend()
|
||||
client = AsyncMock()
|
||||
search = AsyncMock(
|
||||
side_effect=[
|
||||
[3, {b'key-1': {}, b'key-2': {}}],
|
||||
[1, {b'key-3': {}}],
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr(mod, '_DELETE_SCAN_BATCH', 2)
|
||||
monkeypatch.setattr(mod, 'FtSearchLimit', lambda offset, limit: (offset, limit), raising=False)
|
||||
monkeypatch.setattr(mod, 'FtSearchOptions', lambda **kwargs: kwargs, raising=False)
|
||||
monkeypatch.setattr(mod, 'ft', type('FT', (), {'search': search})(), raising=False)
|
||||
|
||||
deleted = await backend._delete_search_results(client, 'idx:col1', '@file_id:{f1}')
|
||||
|
||||
assert deleted == 3
|
||||
assert search.await_count == 2
|
||||
assert [call.args[3]['limit'] for call in search.await_args_list] == [(0, 2), (0, 2)]
|
||||
assert client.delete.await_args_list[0].args == (['key-1', 'key-2'],)
|
||||
assert client.delete.await_args_list[1].args == (['key-3'],)
|
||||
|
||||
async def test_delete_rounds_have_a_hard_stop(self, monkeypatch):
|
||||
mod = get_valkey_module()
|
||||
backend = make_backend()
|
||||
client = AsyncMock()
|
||||
search = AsyncMock(return_value=[2, {b'key': {}}])
|
||||
monkeypatch.setattr(mod, '_DELETE_SCAN_BATCH', 1)
|
||||
monkeypatch.setattr(mod, '_MAX_DELETE_SCAN_ROUNDS', 2)
|
||||
monkeypatch.setattr(mod, 'FtSearchLimit', lambda offset, limit: (offset, limit), raising=False)
|
||||
monkeypatch.setattr(mod, 'FtSearchOptions', lambda **kwargs: kwargs, raising=False)
|
||||
monkeypatch.setattr(mod, 'ft', type('FT', (), {'search': search})(), raising=False)
|
||||
|
||||
with pytest.raises(RuntimeError, match='exceeded 2 batches'):
|
||||
await backend._delete_search_results(client, 'idx:col1', '@file_id:{f1}')
|
||||
|
||||
assert client.delete.await_count == 2
|
||||
|
||||
|
||||
class TestClose:
|
||||
|
||||
Reference in New Issue
Block a user