mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-09 04:40:57 +00:00
feat(cloud): harden multi-tenant runtime resources
This commit is contained in:
Junyan Qin
@@ -6,6 +6,7 @@ Tests the helper methods that don't require real Dify API calls.
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import time
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import langbot_plugin.api.entities.builtin.platform.message as platform_message
|
||||
@@ -18,13 +19,12 @@ class TestDifyWorkflowSubmitClient:
|
||||
|
||||
class FakeResponse:
|
||||
status_code = 503
|
||||
headers = {}
|
||||
|
||||
async def aread(self):
|
||||
return b''
|
||||
|
||||
async def aiter_lines(self):
|
||||
raise AssertionError('error responses must not enter the SSE loop')
|
||||
yield
|
||||
async def aiter_bytes(self, chunk_size=None):
|
||||
del chunk_size
|
||||
if False:
|
||||
yield b''
|
||||
|
||||
class FakeStreamContext:
|
||||
async def __aenter__(self):
|
||||
@@ -66,6 +66,33 @@ class TestDifyWorkflowSubmitClient:
|
||||
)
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sse_parser_rejects_an_unbounded_line(self):
|
||||
from langbot.libs.dify_service_api.v1 import client, errors
|
||||
|
||||
class FakeResponse:
|
||||
async def aiter_bytes(self, chunk_size=None):
|
||||
del chunk_size
|
||||
for _ in range(129):
|
||||
yield b'x' * 8192
|
||||
|
||||
with pytest.raises(errors.DifyAPIError, match='SSE event exceeds'):
|
||||
await anext(client._iter_sse_json(FakeResponse()))
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_rejects_oversized_local_file(self, tmp_path):
|
||||
from langbot.libs.dify_service_api.v1 import client
|
||||
|
||||
file_path = tmp_path / 'large.bin'
|
||||
file_path.write_bytes(b'x' * (client._MAX_DIFY_UPLOAD_BYTES + 1))
|
||||
dify_client = client.AsyncDifyServiceClient(
|
||||
'test-key',
|
||||
'https://dify.example/v1',
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match='exceeds the size limit'):
|
||||
await dify_client.upload_file(file_path, 'person_user-1')
|
||||
|
||||
|
||||
class TestDifyExtractTextOutput:
|
||||
"""Tests for _extract_dify_text_output method."""
|
||||
@@ -320,6 +347,130 @@ class TestDifyHumanInputForms:
|
||||
assert difysvapi._dify_user_from_query(query_a) == difysvapi._dify_user_from_query(query_c)
|
||||
difysvapi._PENDING_FORMS.clear()
|
||||
|
||||
def test_pending_form_lookup_does_not_scan_unrelated_sessions(self, monkeypatch):
|
||||
from langbot.pkg.provider.runners import difysvapi
|
||||
|
||||
def session_key(index: int):
|
||||
return (
|
||||
'instance',
|
||||
f'workspace-{index}',
|
||||
1,
|
||||
'bot',
|
||||
'pipeline',
|
||||
'adapter',
|
||||
'person',
|
||||
f'user-{index}',
|
||||
)
|
||||
|
||||
difysvapi._PENDING_FORMS.clear()
|
||||
for index in range(512):
|
||||
difysvapi._set_pending_form(
|
||||
session_key(index),
|
||||
{
|
||||
'form_token': f'token-{index}',
|
||||
'workflow_run_id': f'run-{index}',
|
||||
},
|
||||
)
|
||||
|
||||
class NoGlobalIterationDict(dict):
|
||||
def __iter__(self):
|
||||
raise AssertionError('pending form lookup scanned all sessions')
|
||||
|
||||
def keys(self):
|
||||
raise AssertionError('pending form lookup scanned all sessions')
|
||||
|
||||
def items(self):
|
||||
raise AssertionError('pending form lookup scanned all sessions')
|
||||
|
||||
def values(self):
|
||||
raise AssertionError('pending form lookup scanned all sessions')
|
||||
|
||||
guarded_forms = NoGlobalIterationDict(difysvapi._PENDING_FORMS)
|
||||
monkeypatch.setattr(difysvapi, '_PENDING_FORMS', guarded_forms)
|
||||
|
||||
assert difysvapi._get_pending_form_by_token(session_key(511), 'token-511')['workflow_run_id'] == 'run-511'
|
||||
difysvapi._set_pending_form(
|
||||
session_key(512),
|
||||
{'form_token': 'token-512', 'workflow_run_id': 'run-512'},
|
||||
)
|
||||
assert len(guarded_forms) == 513
|
||||
|
||||
def test_pending_form_expiry_heap_ignores_stale_overwrite_and_stays_bounded(self):
|
||||
from langbot.pkg.provider.runners import difysvapi
|
||||
|
||||
session_key = (
|
||||
'instance',
|
||||
'workspace',
|
||||
1,
|
||||
'bot',
|
||||
'pipeline',
|
||||
'adapter',
|
||||
'person',
|
||||
'user',
|
||||
)
|
||||
difysvapi._PENDING_FORMS.clear()
|
||||
now = time.time()
|
||||
difysvapi._set_pending_form(
|
||||
session_key,
|
||||
{
|
||||
'form_token': 'token',
|
||||
'workflow_run_id': 'stale',
|
||||
'expiration_time': now + 1,
|
||||
},
|
||||
)
|
||||
for revision in range(500):
|
||||
difysvapi._set_pending_form(
|
||||
session_key,
|
||||
{
|
||||
'form_token': 'token',
|
||||
'workflow_run_id': f'current-{revision}',
|
||||
'expiration_time': now + 3600 + revision,
|
||||
},
|
||||
)
|
||||
|
||||
difysvapi._prune_pending_forms(now + 2)
|
||||
|
||||
assert difysvapi._get_pending_form_by_token(session_key, 'token')['workflow_run_id'] == 'current-499'
|
||||
assert difysvapi._PENDING_FORM_ACTIVE_COUNT == 1
|
||||
assert len(difysvapi._PENDING_FORM_EXPIRY_HEAP) <= max(
|
||||
difysvapi._PENDING_FORM_HEAP_COMPACT_FLOOR,
|
||||
difysvapi._PENDING_FORM_ACTIVE_COUNT * difysvapi._PENDING_FORM_HEAP_MAX_MULTIPLIER,
|
||||
)
|
||||
|
||||
def test_pending_form_capacity_evicts_earliest_session_without_full_scan(
|
||||
self,
|
||||
monkeypatch,
|
||||
):
|
||||
from langbot.pkg.provider.runners import difysvapi
|
||||
|
||||
def session_key(index: int):
|
||||
return (
|
||||
'instance',
|
||||
f'workspace-{index}',
|
||||
1,
|
||||
'bot',
|
||||
'pipeline',
|
||||
'adapter',
|
||||
'person',
|
||||
f'user-{index}',
|
||||
)
|
||||
|
||||
difysvapi._PENDING_FORMS.clear()
|
||||
monkeypatch.setattr(difysvapi, '_PENDING_FORM_MAX_SESSIONS', 2)
|
||||
now = time.time()
|
||||
for index, expires_in in ((1, 300), (2, 100), (3, 200)):
|
||||
difysvapi._set_pending_form(
|
||||
session_key(index),
|
||||
{
|
||||
'form_token': f'token-{index}',
|
||||
'expiration_time': now + expires_in,
|
||||
},
|
||||
)
|
||||
|
||||
assert session_key(1) in difysvapi._PENDING_FORMS
|
||||
assert session_key(2) not in difysvapi._PENDING_FORMS
|
||||
assert session_key(3) in difysvapi._PENDING_FORMS
|
||||
|
||||
def test_interactive_form_data_preserves_pipeline_uuid(self):
|
||||
from langbot.pkg.provider.runners import difysvapi
|
||||
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot.libs.deerflow_api.client import (
|
||||
ERROR_BODY_MAX_BYTES,
|
||||
_read_error_body,
|
||||
)
|
||||
from langbot.libs.deerflow_api.errors import DeerFlowAPIError
|
||||
from langbot.pkg.provider.runners.langflowapi import (
|
||||
_MAX_LANGFLOW_LINE_CHARS,
|
||||
_MAX_LANGFLOW_RESPONSE_BYTES,
|
||||
_iter_limited_lines,
|
||||
_read_limited_response,
|
||||
)
|
||||
|
||||
|
||||
class _ChunkedResponse:
|
||||
def __init__(self, chunks: list[bytes]):
|
||||
self._chunks = chunks
|
||||
|
||||
async def aiter_bytes(self, chunk_size=None):
|
||||
del chunk_size
|
||||
for chunk in self._chunks:
|
||||
yield chunk
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_langflow_rejects_oversized_stream_event():
|
||||
response = _ChunkedResponse([b'x' * (_MAX_LANGFLOW_LINE_CHARS + 1)])
|
||||
|
||||
with pytest.raises(ValueError, match='event exceeds'):
|
||||
await anext(_iter_limited_lines(response))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_langflow_rejects_oversized_blocking_response():
|
||||
response = _ChunkedResponse([b'x' * (_MAX_LANGFLOW_RESPONSE_BYTES + 1)])
|
||||
|
||||
with pytest.raises(ValueError, match='response exceeds'):
|
||||
await _read_limited_response(response)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deerflow_rejects_oversized_error_body():
|
||||
response = _ChunkedResponse([b'x' * (ERROR_BODY_MAX_BYTES + 1)])
|
||||
|
||||
with pytest.raises(DeerFlowAPIError, match='response exceeds'):
|
||||
await _read_error_body(response)
|
||||
@@ -0,0 +1,62 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.provider import runner
|
||||
from langbot.pkg.provider.runners import (
|
||||
cozeapi,
|
||||
dashscopeapi,
|
||||
tboxapi,
|
||||
weknoraapi,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_blocking_provider_iterator_runs_outside_event_loop():
|
||||
release = threading.Event()
|
||||
|
||||
def values():
|
||||
release.wait(timeout=2)
|
||||
yield 'ready'
|
||||
|
||||
task = asyncio.create_task(anext(runner.iterate_sync(values())))
|
||||
await asyncio.sleep(0)
|
||||
assert not task.done()
|
||||
|
||||
release.set()
|
||||
assert await asyncio.wait_for(task, timeout=1) == 'ready'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_provider_iterator_has_event_limit():
|
||||
with pytest.raises(RuntimeError, match='event limit'):
|
||||
async for _ in runner.iterate_sync(iter([1, 2]), max_items=1):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_coze_runner_closes_request_scoped_client():
|
||||
request_runner = object.__new__(cozeapi.CozeAPIRunner)
|
||||
request_runner.coze = AsyncMock()
|
||||
|
||||
await request_runner.aclose()
|
||||
|
||||
request_runner.coze.close.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('append', 'exception_type'),
|
||||
[
|
||||
(cozeapi._append_bounded, ValueError),
|
||||
(dashscopeapi._append_bounded, dashscopeapi.DashscopeAPIError),
|
||||
(tboxapi._append_bounded, tboxapi.TboxAPIError),
|
||||
(weknoraapi._append_bounded, weknoraapi.errors.WeKnoraAPIError),
|
||||
],
|
||||
)
|
||||
def test_provider_accumulators_reject_oversized_output(append, exception_type):
|
||||
with pytest.raises(exception_type, match='exceeds the runtime limit'):
|
||||
append('x' * (1024 * 1024), 'y')
|
||||
@@ -97,7 +97,11 @@ def _query(variables: dict | None = None, context: ExecutionContext = TEST_EXECU
|
||||
|
||||
|
||||
def _register_session(loader: MCPLoader, session: RuntimeMCPSession) -> None:
|
||||
loader.sessions[loader._session_key(session.execution_context, session.server_name)] = session
|
||||
loader._register_session(
|
||||
session.execution_context,
|
||||
session.server_name,
|
||||
session,
|
||||
)
|
||||
|
||||
|
||||
def _http_status_error(status_code: int) -> httpx.HTTPStatusError:
|
||||
@@ -615,3 +619,165 @@ async def test_mcp_loader_shutdown_cancels_startup_tasks_and_closes_sessions_con
|
||||
assert started == {'one', 'two'}
|
||||
assert loader._hosted_mcp_tasks == []
|
||||
assert loader.sessions == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completed_mcp_host_tasks_do_not_accumulate():
|
||||
loader = MCPLoader(_app())
|
||||
task = asyncio.create_task(asyncio.sleep(0))
|
||||
|
||||
loader.track_hosted_task(task, TEST_EXECUTION_CONTEXT)
|
||||
await task
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert loader._hosted_mcp_tasks == []
|
||||
assert loader._hosted_mcp_tasks_by_scope == {}
|
||||
assert loader._scope_generations == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generation_advance_cancels_host_tasks_and_closes_old_sessions():
|
||||
loader = MCPLoader(_app())
|
||||
old_session = SimpleNamespace(
|
||||
server_name='old',
|
||||
shutdown=AsyncMock(),
|
||||
)
|
||||
loader._register_session(
|
||||
TEST_EXECUTION_CONTEXT,
|
||||
old_session.server_name,
|
||||
old_session,
|
||||
)
|
||||
|
||||
async def pending_host():
|
||||
await asyncio.Event().wait()
|
||||
|
||||
hosted_task = asyncio.create_task(pending_host())
|
||||
loader.track_hosted_task(hosted_task, TEST_EXECUTION_CONTEXT)
|
||||
await asyncio.sleep(0)
|
||||
next_context = ExecutionContext(
|
||||
instance_uuid=TEST_EXECUTION_CONTEXT.instance_uuid,
|
||||
workspace_uuid=TEST_EXECUTION_CONTEXT.workspace_uuid,
|
||||
placement_generation=2,
|
||||
)
|
||||
loader.ap.workspace_service.get_execution_binding = AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
instance_uuid=next_context.instance_uuid,
|
||||
workspace_uuid=next_context.workspace_uuid,
|
||||
placement_generation=next_context.placement_generation,
|
||||
)
|
||||
)
|
||||
|
||||
await loader._assert_execution_active(next_context)
|
||||
|
||||
assert hosted_task.cancelled()
|
||||
old_session.shutdown.assert_awaited_once_with()
|
||||
assert loader.sessions == {}
|
||||
assert loader._session_keys_by_scope == {}
|
||||
assert loader._hosted_mcp_tasks_by_scope == {}
|
||||
assert loader._scope_generations == {}
|
||||
|
||||
|
||||
def test_session_lookup_uses_scope_index_without_global_iteration():
|
||||
class NoGlobalIterationDict(dict):
|
||||
def __iter__(self):
|
||||
raise AssertionError('MCP lookup scanned every tenant session')
|
||||
|
||||
def items(self):
|
||||
raise AssertionError('MCP lookup scanned every tenant session')
|
||||
|
||||
def values(self):
|
||||
raise AssertionError('MCP lookup scanned every tenant session')
|
||||
|
||||
loader = MCPLoader(_app())
|
||||
target_context = None
|
||||
target_session = None
|
||||
for index in range(1_000):
|
||||
context = ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid=f'workspace-{index}',
|
||||
placement_generation=1,
|
||||
)
|
||||
session = SimpleNamespace(server_name=f'server-{index}')
|
||||
loader._register_session(context, session.server_name, session)
|
||||
if index == 777:
|
||||
target_context = context
|
||||
target_session = session
|
||||
loader._sessions = NoGlobalIterationDict(loader._sessions)
|
||||
|
||||
assert loader._sessions_for_context(target_context) == [target_session]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_startup_concurrency_is_instance_bounded():
|
||||
app = _app()
|
||||
app.instance_config = SimpleNamespace(data={'mcp': {'lifecycle_concurrency': 2}})
|
||||
loader = MCPLoader(app)
|
||||
active = 0
|
||||
maximum_active = 0
|
||||
release = asyncio.Event()
|
||||
|
||||
async def fake_host(_context, _config):
|
||||
nonlocal active, maximum_active
|
||||
active += 1
|
||||
maximum_active = max(maximum_active, active)
|
||||
if maximum_active == 2:
|
||||
release.set()
|
||||
await release.wait()
|
||||
await asyncio.sleep(0)
|
||||
active -= 1
|
||||
|
||||
loader._host_mcp_server = fake_host
|
||||
|
||||
await asyncio.gather(
|
||||
*(
|
||||
loader.host_mcp_server(
|
||||
TEST_EXECUTION_CONTEXT,
|
||||
{'name': f'server-{index}'},
|
||||
)
|
||||
for index in range(20)
|
||||
)
|
||||
)
|
||||
|
||||
assert loader._lifecycle_concurrency == 2
|
||||
assert maximum_active == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_startup_dispatcher_does_not_create_every_server_task_at_once():
|
||||
app = _app()
|
||||
app.instance_config = SimpleNamespace(data={'mcp': {'lifecycle_concurrency': 2}})
|
||||
loader = MCPLoader(app)
|
||||
started = 0
|
||||
first_batch_started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def fake_host(_context, _config):
|
||||
nonlocal started
|
||||
started += 1
|
||||
if started == 2:
|
||||
first_batch_started.set()
|
||||
await release.wait()
|
||||
|
||||
loader.host_mcp_server = fake_host
|
||||
configs = [(TEST_EXECUTION_CONTEXT, {'name': f'server-{index}'}) for index in range(20)]
|
||||
|
||||
dispatch_task = asyncio.create_task(loader._host_server_configs_bounded(configs))
|
||||
await asyncio.wait_for(first_batch_started.wait(), timeout=1)
|
||||
|
||||
assert started == 2
|
||||
assert len(loader._hosted_mcp_tasks) == 2
|
||||
|
||||
release.set()
|
||||
await dispatch_task
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert started == 20
|
||||
assert loader._hosted_mcp_tasks == []
|
||||
assert loader._hosted_mcp_tasks_by_scope == {}
|
||||
|
||||
|
||||
def test_invalid_mcp_lifecycle_concurrency_uses_safe_default():
|
||||
app = _app()
|
||||
app.instance_config = SimpleNamespace(data={'mcp': {'lifecycle_concurrency': True}})
|
||||
|
||||
assert MCPLoader(app)._lifecycle_concurrency == 16
|
||||
|
||||
@@ -7,6 +7,7 @@ and error handling without calling real LLM APIs.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import pytest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
@@ -18,7 +19,7 @@ from langbot.pkg.entity.errors import provider as provider_errors
|
||||
from langbot.pkg.provider.modelmgr import token
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
from langbot.pkg.workspace.entities import WorkspaceExecutionBinding
|
||||
from langbot.pkg.workspace.errors import WorkspaceGenerationMismatchError
|
||||
from langbot.pkg.workspace.errors import WorkspaceGenerationMismatchError, WorkspaceInvariantError
|
||||
from tests.unit_tests.provider.conftest import (
|
||||
TEST_EXECUTION_CONTEXT,
|
||||
TEST_WORKSPACE_UUID,
|
||||
@@ -125,6 +126,24 @@ async def test_model_manager_load_models_from_db(fake_requester_registry, fake_p
|
||||
assert len(model_mgr.rerank_model_dict) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_cloud_workspace_does_not_retain_generation(
|
||||
mock_app_for_modelmgr,
|
||||
):
|
||||
model_mgr = ModelManager(mock_app_for_modelmgr)
|
||||
|
||||
await model_mgr._load_workspace_models(TEST_EXECUTION_CONTEXT)
|
||||
|
||||
assert model_mgr.provider_dict == {}
|
||||
assert model_mgr.llm_model_dict == {}
|
||||
assert model_mgr.embedding_model_dict == {}
|
||||
assert model_mgr.rerank_model_dict == {}
|
||||
assert model_mgr._scope_generations == {}
|
||||
|
||||
await model_mgr.resolve_execution_context(TEST_EXECUTION_CONTEXT)
|
||||
assert model_mgr._scope_generations == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_manager_load_provider_unknown_requester(mock_app_for_modelmgr):
|
||||
"""Test ModelManager raises RequesterNotFoundError for unknown requester."""
|
||||
@@ -939,6 +958,110 @@ async def test_runtime_cache_rejects_stale_placement_generation(fake_requester_r
|
||||
await model_mgr.get_model_by_uuid(stale_context, 'any-model')
|
||||
|
||||
|
||||
def test_generation_advance_prunes_superseded_model_runtime_objects():
|
||||
class NoGlobalIterationDict(dict):
|
||||
def __iter__(self):
|
||||
raise AssertionError('generation advance scanned every model runtime')
|
||||
|
||||
def items(self):
|
||||
raise AssertionError('generation advance scanned every model runtime')
|
||||
|
||||
def keys(self):
|
||||
raise AssertionError('generation advance scanned every model runtime')
|
||||
|
||||
model_mgr = ModelManager(Mock())
|
||||
old_context = ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=1,
|
||||
)
|
||||
new_context = dataclasses.replace(old_context, placement_generation=2)
|
||||
model_mgr._observe_execution_context(old_context)
|
||||
for cache in (
|
||||
model_mgr.provider_dict,
|
||||
model_mgr.llm_model_dict,
|
||||
model_mgr.embedding_model_dict,
|
||||
model_mgr.rerank_model_dict,
|
||||
):
|
||||
model_mgr._cache_set(
|
||||
cache,
|
||||
('instance-a', 'workspace-a', 1, 'resource-a'),
|
||||
object(),
|
||||
)
|
||||
model_mgr._cache_set(
|
||||
cache,
|
||||
('instance-a', 'workspace-b', 1, 'resource-b'),
|
||||
object(),
|
||||
)
|
||||
|
||||
model_mgr.provider_dict = NoGlobalIterationDict(model_mgr.provider_dict)
|
||||
model_mgr.llm_model_dict = NoGlobalIterationDict(model_mgr.llm_model_dict)
|
||||
model_mgr.embedding_model_dict = NoGlobalIterationDict(model_mgr.embedding_model_dict)
|
||||
model_mgr.rerank_model_dict = NoGlobalIterationDict(model_mgr.rerank_model_dict)
|
||||
|
||||
model_mgr._observe_execution_context(new_context)
|
||||
|
||||
for cache in (
|
||||
model_mgr.provider_dict,
|
||||
model_mgr.llm_model_dict,
|
||||
model_mgr.embedding_model_dict,
|
||||
model_mgr.rerank_model_dict,
|
||||
):
|
||||
assert ('instance-a', 'workspace-a', 1, 'resource-a') not in cache
|
||||
assert ('instance-a', 'workspace-b', 1, 'resource-b') in cache
|
||||
with pytest.raises(WorkspaceInvariantError, match='rolled back'):
|
||||
model_mgr._observe_execution_context(old_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generation_advance_closes_retired_provider_requester(
|
||||
fake_requester_registry,
|
||||
runtime_provider,
|
||||
):
|
||||
model_mgr = fake_requester_registry
|
||||
runtime_provider.requester.aclose = AsyncMock()
|
||||
await model_mgr.cache_provider(TEST_EXECUTION_CONTEXT, runtime_provider)
|
||||
|
||||
next_context = dataclasses.replace(
|
||||
TEST_EXECUTION_CONTEXT,
|
||||
placement_generation=2,
|
||||
)
|
||||
model_mgr.ap.workspace_service.get_execution_binding = AsyncMock(
|
||||
return_value=WorkspaceExecutionBinding(
|
||||
instance_uuid=next_context.instance_uuid,
|
||||
workspace_uuid=next_context.workspace_uuid,
|
||||
placement_generation=next_context.placement_generation,
|
||||
write_fenced=False,
|
||||
state='active',
|
||||
)
|
||||
)
|
||||
|
||||
await model_mgr.resolve_execution_context(next_context)
|
||||
|
||||
runtime_provider.requester.aclose.assert_awaited_once_with()
|
||||
assert model_mgr.provider_dict == {}
|
||||
assert model_mgr._scope_generations == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_manager_shutdown_closes_all_requesters_once(
|
||||
fake_requester_registry,
|
||||
runtime_provider,
|
||||
):
|
||||
model_mgr = fake_requester_registry
|
||||
runtime_provider.requester.aclose = AsyncMock()
|
||||
await model_mgr.cache_provider(TEST_EXECUTION_CONTEXT, runtime_provider)
|
||||
|
||||
await model_mgr.shutdown()
|
||||
await model_mgr.shutdown()
|
||||
|
||||
runtime_provider.requester.aclose.assert_awaited_once_with()
|
||||
assert model_mgr.provider_dict == {}
|
||||
assert model_mgr.llm_model_dict == {}
|
||||
assert model_mgr.embedding_model_dict == {}
|
||||
assert model_mgr.rerank_model_dict == {}
|
||||
|
||||
|
||||
def test_provider_not_found_error_str():
|
||||
"""Test ProviderNotFoundError string representation."""
|
||||
error = provider_errors.ProviderNotFoundError('test-provider')
|
||||
|
||||
@@ -16,6 +16,8 @@ from importlib import import_module
|
||||
|
||||
import langbot_plugin.api.entities.builtin.provider.session as provider_session
|
||||
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
||||
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
||||
import langbot_plugin.api.entities.builtin.provider.prompt as provider_prompt
|
||||
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
from langbot.pkg.pipeline.pool import (
|
||||
@@ -426,3 +428,131 @@ class TestSessionManagerWorkspaceIsolation:
|
||||
|
||||
with pytest.raises(ExecutionContextMismatchError):
|
||||
await manager.get_conversation(query, session, [], 'pipeline-1', 'bot-b')
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_per_workspace_capacity_evicts_oldest_idle_session(self):
|
||||
manager = self.manager()
|
||||
manager.ap.instance_config.data['system'] = {
|
||||
'session_retention': {
|
||||
'max_entries': 10,
|
||||
'max_entries_per_workspace': 2,
|
||||
}
|
||||
}
|
||||
queries = []
|
||||
sessions = []
|
||||
for index in range(3):
|
||||
query = scoped_query()
|
||||
query.launcher_id = f'launcher-{index}'
|
||||
queries.append(query)
|
||||
sessions.append(await manager.get_session(query))
|
||||
|
||||
assert len(manager.session_list) == 2
|
||||
assert sessions[0] not in manager.session_list
|
||||
assert sessions[1:] == manager.session_list
|
||||
assert await manager.get_session(queries[-1]) is manager.session_list[-1]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_new_session_does_not_scan_other_workspace_sessions(self):
|
||||
manager = self.manager()
|
||||
manager.ap.instance_config.data['system'] = {
|
||||
'session_retention': {
|
||||
'max_entries': 600,
|
||||
'max_entries_per_workspace': 2,
|
||||
}
|
||||
}
|
||||
for index in range(512):
|
||||
query = scoped_query(workspace_uuid=f'workspace-{index}')
|
||||
query.launcher_id = f'launcher-{index}'
|
||||
await manager.get_session(query)
|
||||
|
||||
class NoGlobalIterationDict(dict):
|
||||
def __iter__(self):
|
||||
raise AssertionError('global session index iteration is forbidden')
|
||||
|
||||
def items(self):
|
||||
raise AssertionError('global session index iteration is forbidden')
|
||||
|
||||
def values(self):
|
||||
raise AssertionError('global session index iteration is forbidden')
|
||||
|
||||
manager._session_index = NoGlobalIterationDict(manager._session_index)
|
||||
query = scoped_query(workspace_uuid='workspace-new')
|
||||
query.launcher_id = 'launcher-new'
|
||||
|
||||
session = await manager.get_session(query)
|
||||
|
||||
assert session.workspace_uuid == 'workspace-new'
|
||||
assert len(manager._session_index) == 513
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stale_expiry_revision_does_not_evict_recent_session(
|
||||
self,
|
||||
monkeypatch,
|
||||
):
|
||||
sessionmgr = get_session_module()
|
||||
manager = self.manager()
|
||||
manager.ap.instance_config.data['system'] = {
|
||||
'session_retention': {
|
||||
'max_entries': 10,
|
||||
'max_entries_per_workspace': 10,
|
||||
'idle_ttl_seconds': 1,
|
||||
}
|
||||
}
|
||||
clock = [0.0]
|
||||
monkeypatch.setattr(sessionmgr.time, 'monotonic', lambda: clock[0])
|
||||
first_query = scoped_query()
|
||||
first_query.launcher_id = 'first'
|
||||
first = await manager.get_session(first_query)
|
||||
|
||||
clock[0] = 0.5
|
||||
assert await manager.get_session(first_query) is first
|
||||
|
||||
clock[0] = 1.25
|
||||
second_query = scoped_query()
|
||||
second_query.launcher_id = 'second'
|
||||
await manager.get_session(second_query)
|
||||
assert first in manager.session_list
|
||||
|
||||
clock[0] = 2.0
|
||||
third_query = scoped_query()
|
||||
third_query.launcher_id = 'third'
|
||||
await manager.get_session(third_query)
|
||||
assert first not in manager.session_list
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_access_revision_heap_stays_bounded(self):
|
||||
manager = self.manager()
|
||||
query = scoped_query()
|
||||
await manager.get_session(query)
|
||||
|
||||
for _ in range(1000):
|
||||
await manager.get_session(query)
|
||||
|
||||
assert len(manager._session_expiry_heap) <= 64
|
||||
|
||||
def test_trim_conversation_drops_retained_binary_payloads(self):
|
||||
manager = self.manager()
|
||||
conversation = provider_session.Conversation(
|
||||
prompt=provider_prompt.Prompt(name='test', messages=[]),
|
||||
messages=[
|
||||
provider_message.Message(
|
||||
role='user',
|
||||
content=[
|
||||
provider_message.ContentElement.from_text('hello'),
|
||||
provider_message.ContentElement.from_image_base64('x' * 1000000),
|
||||
provider_message.ContentElement.from_file_base64(
|
||||
'y' * 1000000,
|
||||
'large.bin',
|
||||
),
|
||||
],
|
||||
)
|
||||
],
|
||||
pipeline_uuid='pipeline-1',
|
||||
bot_uuid=TEST_BOT_UUID,
|
||||
)
|
||||
|
||||
manager.trim_conversation_messages(conversation, max_rounds=10)
|
||||
|
||||
content = conversation.messages[0].content
|
||||
assert content[1].image_base64 is None
|
||||
assert content[2].file_base64 is None
|
||||
|
||||
@@ -4,6 +4,7 @@ import base64
|
||||
import contextlib
|
||||
import os
|
||||
import tempfile
|
||||
import threading
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
@@ -459,6 +460,24 @@ async def test_edit_rejects_missing_string():
|
||||
assert 'not found' in result['error'].lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_rejects_oversized_host_file(monkeypatch):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
loader, _ = _make_loader_with_workspace(tmpdir)
|
||||
with open(os.path.join(tmpdir, 'large.txt'), 'wb') as f:
|
||||
f.write(b'12345')
|
||||
monkeypatch.setattr(native_loader, '_MAX_HOST_EDIT_FILE_BYTES', 4)
|
||||
|
||||
result = await loader.invoke_tool(
|
||||
'edit',
|
||||
{'path': '/workspace/large.txt', 'old_string': '1', 'new_string': 'x'},
|
||||
_make_query(),
|
||||
)
|
||||
|
||||
assert result['ok'] is False
|
||||
assert 'edit limit' in result['error']
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_path_escape_blocked():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
@@ -678,6 +697,34 @@ async def test_glob_caps_match_count_and_returns_preview():
|
||||
assert result['truncated_by'] == 'matches'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_glob_runs_off_event_loop_and_caps_directory_walk(monkeypatch):
|
||||
monkeypatch.setattr(native_loader, '_FILE_WALK_MAX_ENTRIES', 10)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
loader, _ = _make_loader_with_workspace(tmpdir)
|
||||
event_loop_thread = threading.get_ident()
|
||||
observed_threads: list[int] = []
|
||||
original = loader._glob_host_location
|
||||
|
||||
def observe(*args, **kwargs):
|
||||
observed_threads.append(threading.get_ident())
|
||||
return original(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(loader, '_glob_host_location', observe)
|
||||
for index in range(12):
|
||||
with open(os.path.join(tmpdir, f'file-{index:03d}.txt'), 'w', encoding='utf-8') as f:
|
||||
f.write(str(index))
|
||||
|
||||
result = await loader.invoke_tool('glob', {'path': '/workspace', 'pattern': '*.txt'}, _make_query())
|
||||
|
||||
assert result['ok'] is True
|
||||
assert result['total'] == 10
|
||||
assert result['truncated'] is True
|
||||
assert result['truncated_by'] == 'scan'
|
||||
assert observed_threads and observed_threads[0] != event_loop_thread
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_grep_reports_invalid_regex_and_truncates_long_matching_lines():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
@@ -695,3 +742,21 @@ async def test_grep_reports_invalid_regex_and_truncates_long_matching_lines():
|
||||
assert result['truncated_by'] == 'line'
|
||||
assert result['matches'][0]['file'] == '/workspace/data.txt'
|
||||
assert result['matches'][0]['content'].endswith('... [truncated]')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_grep_interrupts_catastrophic_regex(monkeypatch):
|
||||
monkeypatch.setattr(native_loader, '_GREP_REGEX_TIMEOUT_SECONDS', 0.001)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
loader, _ = _make_loader_with_workspace(tmpdir)
|
||||
with open(os.path.join(tmpdir, 'data.txt'), 'w', encoding='utf-8') as f:
|
||||
f.write(('a' * 100_000) + '!')
|
||||
|
||||
result = await loader.invoke_tool(
|
||||
'grep',
|
||||
{'path': '/workspace', 'pattern': r'(a+)+$'},
|
||||
_make_query(),
|
||||
)
|
||||
|
||||
assert result == {'ok': False, 'error': 'Regex search timed out'}
|
||||
|
||||
Reference in New Issue
Block a user