xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-08-09 12:40:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(tenancy): add Workspace multi-tenant foundation (#2353)

Browse Source 
* Document multi-tenant workspace architecture

* Add OSS and commercial workspace boundaries

* docs: redesign multi-tenant workspace architecture

* feat(tenancy): implement workspace isolation

* docs(tenancy): record verification evidence

* docs(tenancy): revise single-instance SaaS topology

* docs(tenancy): refine architecture options

* docs: finalize cloud v2 multi-tenant decisions

* feat(tenancy): establish cloud isolation foundations

* feat(tenancy): harden shared cloud runtime boundaries

* docs(tenancy): record final isolation verification

* fix(tenancy): close isolation and permission gaps

* docs(tenancy): record final isolation verification

* feat(tenancy): connect cloud workspace control plane

* fix(build): install git for pinned SDK

* docs(cloud): update control plane verification

* chore: update multi-tenant SDK pin

* fix(cloud): skip legacy model sync during startup

* test(cloud): preserve minimal model manager fixtures

* fix(cloud): preserve authenticated account context

* fix(cloud): reuse authenticated account for user info

* feat(cloud): complete Workspace settings navigation

* test(web): cover Workspace dropdown menu

* feat(web): place workspace controls in sidebar

* refactor(web): streamline workspace controls

* style(web): format workspace layout test

* fix(cloud): surface runtime and workspace plan status

* fix(plugin): keep runtime identity stable across restarts

* fix(ui): widen and center workspace switcher

* fix(ui): hide roles from workspace switcher

* fix(ui): align workspace switcher with sidebar entries

* feat(workspace): add in-product collaboration and direct Cloud launch

* style: format collaboration changes

* fix(workspace): bind collaboration APIs to tenant UoW

* fix(cloud): preserve Core-owned collaboration state

* test(cloud): require Space identity for invite registration

* feat(cloud): complete secure invitation experience

* style(web): format invitation flows

* fix(cloud): recover box runtime without unscoped skill reload

* feat(oss): enforce invitation account and owner billing flows

* style: format OSS account service

* test(oss): cover invitation logout handoff

* fix(oss): resolve workspace owner in scoped session

* feat(cloud): harden multi-tenant runtime resources

* fix(cloud): bound runtime restart storms

* fix(cloud): eliminate periodic runtime CPU spikes

* fix(cloud): enforce instance capacity ceilings

* fix(cloud): scope public login capability discovery

* fix(cloud): bound tenant maintenance and monitoring work

* fix(runtime): bound tenant resource amplification

* fix(deps): pin green multi-tenant plugin SDK

* fix(cloud): handle unavailable skill capability

* fix(security): require authentication for image file endpoint (H-2)

- Changed /api/v1/files/image from AuthType.NONE to USER_TOKEN_OR_API_KEY
- Added Permission.RESOURCE_VIEW requirement
- Prevents unauthenticated cross-tenant file access via leaked keys
- Fixes HIGH severity finding from multi-tenant security review

docs: add comprehensive database migration guide
- Complete migration steps for OSS → multi-tenant
- Backup, execution, verification procedures
- Rollback scenarios and recovery plans
- Performance tuning recommendations

* test: add comprehensive cross-tenant isolation tests

Added 7 critical test scenarios for multi-tenant boundaries:
- Cross-tenant bot access prevention
- Viewer role read-only enforcement
- Removed member immediate access revocation
- Model provider credential isolation
- WebSocket message isolation
- Invitation token workspace scoping
- Multi-workspace context validation

These tests address P0-2 coverage gaps for:
- workspaces.py (membership & invitation flows)
- user.py (authentication & authorization)
- websocket_chat.py (real-time isolation)
- plugins.py (resource access control)

docs: finalize database migration guide

* fix(security): resolve M-1, M-2, M-3 security findings

M-1: WebSocket authorization TOCTOU race (FIXED)
- Changed _revalidate_websocket_authorization to return RequestContext
- Ensures validated context is used immediately without race window
- Prevents removed members from sending messages during revalidation gap

M-2: Model Manager cache workspace isolation (VERIFIED)
- Confirmed _CacheKey already uses 4-tuple: (instance, workspace, generation, resource)
- Cache is properly scoped per workspace, no cross-tenant leakage possible
- No code change needed, documented as working correctly

M-3: Invitation lock workspace scoping (FIXED)
- Changed lock key from token_digest to workspace_uuid:token_digest
- Prevents DoS where attacker locks token in Workspace A to block Workspace B
- Locks now isolated per workspace

All MEDIUM severity findings from security review now resolved.

* fix(cloud): unblock tenant CI and enforce knowledge quotas

* fix(tenancy): scope rerank model sync

---------

Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>
This commit is contained in:
RockChinQ
2026-07-30 21:43:35 +08:00
committed by GitHub
parent 463b120923
commit e1ac5e0fc8
468 changed files with 78320 additions and 13137 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unit_tests/provider/conftest.py
+43
View File
@@ -16,6 +16,18 @@ from langbot.pkg.provider.modelmgr import token
from langbot.pkg.provider.modelmgr.modelmgr import ModelManager
from langbot.pkg.entity.persistence import model as persistence_model
from langbot.pkg.discover import engine as discover_engine
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.workspace.entities import WorkspaceExecutionBinding
TEST_INSTANCE_UUID = 'test-instance'
TEST_WORKSPACE_UUID = 'test-workspace'
TEST_GENERATION = 1
TEST_EXECUTION_CONTEXT = ExecutionContext(
instance_uuid=TEST_INSTANCE_UUID,
workspace_uuid=TEST_WORKSPACE_UUID,
placement_generation=TEST_GENERATION,
)
class FakeProviderAPIRequester(requester.ProviderAPIRequester):
@@ -157,6 +169,26 @@ def mock_app_for_modelmgr():
app.llm_model_service = AsyncMock()
app.embedding_models_service = AsyncMock()
app.monitoring_service = AsyncMock()
app.workspace_service = SimpleNamespace(
get_execution_binding=AsyncMock(
return_value=WorkspaceExecutionBinding(
instance_uuid=TEST_INSTANCE_UUID,
workspace_uuid=TEST_WORKSPACE_UUID,
placement_generation=TEST_GENERATION,
write_fenced=False,
state='active',
)
),
get_local_execution_binding=AsyncMock(
return_value=WorkspaceExecutionBinding(
instance_uuid=TEST_INSTANCE_UUID,
workspace_uuid=TEST_WORKSPACE_UUID,
placement_generation=TEST_GENERATION,
write_fenced=False,
state='active',
)
),
)
return app
@@ -184,6 +216,7 @@ def fake_persistence_data():
providers = [
persistence_model.ModelProvider(
workspace_uuid=TEST_WORKSPACE_UUID,
uuid=provider_uuid,
name='Test Provider',
requester='fake-requester',
@@ -191,6 +224,7 @@ def fake_persistence_data():
api_keys=['test-api-key-1', 'test-api-key-2'],
),
persistence_model.ModelProvider(
workspace_uuid=TEST_WORKSPACE_UUID,
uuid=provider_uuid2,
name='Test Provider 2',
requester='another-fake-requester',
@@ -201,6 +235,7 @@ def fake_persistence_data():
llm_models = [
persistence_model.LLMModel(
workspace_uuid=TEST_WORKSPACE_UUID,
uuid='test-llm-uuid-1',
name='TestLLM-1',
provider_uuid=provider_uuid,
@@ -208,6 +243,7 @@ def fake_persistence_data():
extra_args={'temperature': 0.7},
),
persistence_model.LLMModel(
workspace_uuid=TEST_WORKSPACE_UUID,
uuid='test-llm-uuid-2',
name='TestLLM-2',
provider_uuid=provider_uuid,
@@ -218,6 +254,7 @@ def fake_persistence_data():
embedding_models = [
persistence_model.EmbeddingModel(
workspace_uuid=TEST_WORKSPACE_UUID,
uuid='test-embedding-uuid-1',
name='TestEmbedding-1',
provider_uuid=provider_uuid,
@@ -227,6 +264,7 @@ def fake_persistence_data():
rerank_models = [
persistence_model.RerankModel(
workspace_uuid=TEST_WORKSPACE_UUID,
uuid='test-rerank-uuid-1',
name='TestRerank-1',
provider_uuid=provider_uuid2,
@@ -252,6 +290,7 @@ def runtime_provider(fake_persistence_data, mock_app_for_modelmgr):
requester_inst = FakeProviderAPIRequester(mock_app_for_modelmgr, {'base_url': provider_entity.base_url})
return requester.RuntimeProvider(
execution_context=TEST_EXECUTION_CONTEXT,
provider_entity=provider_entity,
token_mgr=token_mgr,
requester=requester_inst,
@@ -263,6 +302,7 @@ def runtime_llm_model(fake_persistence_data, runtime_provider):
"""Provides a RuntimeLLMModel instance for testing."""
model_entity = fake_persistence_data['llm_models'][0]
return requester.RuntimeLLMModel(
execution_context=TEST_EXECUTION_CONTEXT,
model_entity=model_entity,
provider=runtime_provider,
)
@@ -273,6 +313,7 @@ def runtime_embedding_model(fake_persistence_data, runtime_provider):
"""Provides a RuntimeEmbeddingModel instance for testing."""
model_entity = fake_persistence_data['embedding_models'][0]
return requester.RuntimeEmbeddingModel(
execution_context=TEST_EXECUTION_CONTEXT,
model_entity=model_entity,
provider=runtime_provider,
)
@@ -286,6 +327,7 @@ def runtime_rerank_model(fake_persistence_data, mock_app_for_modelmgr):
requester_inst = AnotherFakeRequester(mock_app_for_modelmgr, {'base_url': provider_entity.base_url})
provider = requester.RuntimeProvider(
execution_context=TEST_EXECUTION_CONTEXT,
provider_entity=provider_entity,
token_mgr=token_mgr,
requester=requester_inst,
@@ -293,6 +335,7 @@ def runtime_rerank_model(fake_persistence_data, mock_app_for_modelmgr):
model_entity = fake_persistence_data['rerank_models'][0]
return requester.RuntimeRerankModel(
execution_context=TEST_EXECUTION_CONTEXT,
model_entity=model_entity,
provider=provider,
)
tests/unit_tests/provider/runners/test_difysvapi_runner.py
+181 -7
View File
@@ -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."""
@@ -252,42 +279,65 @@ class TestDifyHumanInputForms:
runner.dify_client.upload_file = AsyncMock(return_value={'id': 'upload-1'})
return runner
def test_pending_forms_are_isolated_by_bot_and_pipeline(self):
def test_pending_forms_are_isolated_by_workspace_generation_bot_and_pipeline(self):
from langbot.pkg.provider.runners import difysvapi
query_a = MagicMock()
query_a.instance_uuid = 'instance-a'
query_a.workspace_uuid = 'workspace-a'
query_a.placement_generation = 1
query_a.bot_uuid = 'bot-a'
query_a.pipeline_uuid = 'pipeline-a'
query_a.session.launcher_type.value = 'person'
query_a.session.launcher_id = 'shared-user'
query_b = MagicMock()
query_b.instance_uuid = 'instance-a'
query_b.workspace_uuid = 'workspace-a'
query_b.placement_generation = 1
query_b.bot_uuid = 'bot-b'
query_b.pipeline_uuid = 'pipeline-a'
query_b.session.launcher_type.value = 'person'
query_b.session.launcher_id = 'shared-user'
query_c = MagicMock()
query_c.instance_uuid = 'instance-a'
query_c.workspace_uuid = 'workspace-a'
query_c.placement_generation = 1
query_c.bot_uuid = 'bot-a'
query_c.pipeline_uuid = 'pipeline-b'
query_c.session.launcher_type.value = 'person'
query_c.session.launcher_id = 'shared-user'
query_d = MagicMock()
query_d.instance_uuid = 'instance-a'
query_d.workspace_uuid = 'workspace-b'
query_d.placement_generation = 2
query_d.bot_uuid = 'bot-a'
query_d.pipeline_uuid = 'pipeline-a'
query_d.session.launcher_type.value = 'person'
query_d.session.launcher_id = 'shared-user'
key_a = difysvapi._session_key_from_query(query_a)
key_b = difysvapi._session_key_from_query(query_b)
key_c = difysvapi._session_key_from_query(query_c)
key_d = difysvapi._session_key_from_query(query_d)
difysvapi._PENDING_FORMS.clear()
difysvapi._set_pending_form(key_a, {'form_token': 'token-a', 'workflow_run_id': 'run-a'})
difysvapi._set_pending_form(key_b, {'form_token': 'token-b', 'workflow_run_id': 'run-b'})
difysvapi._set_pending_form(key_c, {'form_token': 'token-c', 'workflow_run_id': 'run-c'})
difysvapi._set_pending_form(key_d, {'form_token': 'token-d', 'workflow_run_id': 'run-d'})
assert key_a != key_b
assert key_a != key_c
assert key_a != key_d
assert difysvapi._get_pending_form_by_token(key_a, 'token-a') is not None
assert difysvapi._get_pending_form_by_token(key_a, 'token-b') is None
assert difysvapi._get_pending_form_by_token(key_a, 'token-c') is None
assert difysvapi._get_pending_form_by_token(key_a, 'token-d') is None
assert difysvapi._get_pending_form_by_token(key_b, 'token-b') is not None
assert difysvapi._get_pending_form_by_token(key_c, 'token-c') is not None
assert difysvapi._get_pending_form_by_token(key_d, 'token-d') is not None
assert difysvapi._get_latest_pending_form(key_a)['workflow_run_id'] == 'run-a'
assert difysvapi._get_latest_pending_form(key_b)['workflow_run_id'] == 'run-b'
assert difysvapi._get_latest_pending_form(key_c)['workflow_run_id'] == 'run-c'
@@ -297,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
tests/unit_tests/provider/runners/test_remote_stream_limits.py
+49
View File
@@ -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)
tests/unit_tests/provider/runners/test_runner_resource_limits.py
+62
View File
@@ -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')
tests/unit_tests/provider/test_localagent_no_duplicate.py
+13 -1
View File
@@ -10,6 +10,7 @@ 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.session as provider_session
from langbot.pkg.api.http.context import ExecutionContext, PrincipalContext, PrincipalType
from langbot.pkg.provider.runners.localagent import LocalAgentRunner
@@ -97,7 +98,7 @@ def make_query() -> pipeline_query.Query:
adapter = AsyncMock()
adapter.is_stream_output_supported = AsyncMock(return_value=False)
return pipeline_query.Query.model_construct(
query = pipeline_query.Query.model_construct(
query_id='no-dup-query',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
@@ -124,6 +125,17 @@ def make_query() -> pipeline_query.Query:
use_llm_model_uuid='test-model-uuid',
variables={},
)
object.__setattr__(
query,
'_execution_context',
ExecutionContext(
instance_uuid='instance-test',
workspace_uuid='workspace-test',
placement_generation=1,
trigger_principal=PrincipalContext(PrincipalType.SYSTEM),
),
)
return query
def _make_app(provider) -> SimpleNamespace:
tests/unit_tests/provider/test_localagent_sandbox_exec.py
+16 -1
View File
@@ -10,6 +10,7 @@ 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.session as provider_session
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.provider.runners.localagent import LocalAgentRunner, _StreamAccumulator
@@ -95,7 +96,7 @@ def make_query() -> pipeline_query.Query:
adapter = AsyncMock()
adapter.is_stream_output_supported = AsyncMock(return_value=False)
return pipeline_query.Query.model_construct(
query = pipeline_query.Query.model_construct(
query_id='avg-query',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
@@ -122,6 +123,20 @@ def make_query() -> pipeline_query.Query:
use_llm_model_uuid='test-model-uuid',
variables={},
)
object.__setattr__(
query,
'_execution_context',
ExecutionContext(
instance_uuid='instance-test',
workspace_uuid='workspace-test',
placement_generation=1,
bot_uuid='bot-uuid',
pipeline_uuid='pipeline-uuid',
query_uuid='query-avg',
),
)
object.__setattr__(query, 'query_uuid', 'query-avg')
return query
def test_stream_accumulator_merges_fragmented_tool_call_arguments():
tests/unit_tests/provider/test_mcp_box_integration.py
+96 -11
View File
@@ -154,7 +154,17 @@ def mcp_module():
def _make_ap():
ap = Mock()
ap.logger = Mock()
ap.instance_config = SimpleNamespace(data={'mcp': {'stdio': {'enabled': True}}})
ap.workspace_service = Mock()
ap.workspace_service.get_execution_binding = AsyncMock(
return_value=SimpleNamespace(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
)
ap.box_service = Mock()
ap.box_service.get_managed_process_websocket_connection = AsyncMock(return_value=('ws://box.example/process', {}))
return ap
@@ -166,6 +176,11 @@ def _make_session(mcp_module, server_config: dict, ap=None):
server_config=server_config,
enable=True,
ap=ap,
execution_context=mcp_module.ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
),
)
@@ -591,6 +606,26 @@ class TestGetRuntimeInfoDict:
assert info['status'] == 'connecting'
assert 'box_session_id' not in info
def test_runtime_error_detail_never_echoes_secret_config(self, mcp_module):
s = _make_session(
mcp_module,
{
'name': 'test',
'uuid': 'test-uuid',
'mode': 'invalid',
'headers': {'Authorization': 'Bearer TOPSECRET'},
'env': {'API_KEY': 'TOPSECRET'},
},
)
s.status = mcp_module.MCPSessionStatus.ERROR
s.error_message = f'Unknown MCP server mode: {s.server_config}'
info = s.get_runtime_info_dict()
assert info['error_message'] == 'MCP runtime failed'
assert info['error_code'] == 'runtime_error'
assert 'TOPSECRET' not in str(info)
def test_runtime_tools_include_parameters(self, mcp_module):
s = _make_session(
mcp_module,
@@ -794,7 +829,7 @@ class TestGetRuntimeInfoDict:
assert ap.box_service.available is True
@pytest.mark.asyncio
async def test_enabled_box_timeout_does_not_exhaust_mcp_retry_budget(self, mcp_module, monkeypatch):
async def test_enabled_box_timeout_does_not_exhaust_mcp_retry_budget(self, mcp_module):
ap = _make_ap()
ap.box_service.available = False
ap.box_service.enabled = True
@@ -819,14 +854,13 @@ class TestGetRuntimeInfoDict:
raise RuntimeError('Box runtime is not available after 1 seconds')
session._lifecycle_loop = lifecycle
sleep = AsyncMock()
monkeypatch.setattr(mcp_module.asyncio, 'sleep', sleep)
session._sleep_with_execution_fence = AsyncMock()
await session._lifecycle_loop_with_retry()
assert attempts == 2
assert session.retry_count == 0
sleep.assert_awaited_once_with(1)
session._sleep_with_execution_fence.assert_awaited_once_with(1)
@pytest.mark.asyncio
async def test_disabled_box_still_stops_mcp_retry_loop(self, mcp_module):
@@ -914,6 +948,31 @@ class TestBoxConfigParsing:
assert s.box_config.host_path_mode == 'ro'
@pytest.mark.asyncio
async def test_stdio_instance_gate_runs_before_box_transport(mcp_module):
ap = _make_ap()
ap.instance_config.data['mcp']['stdio']['enabled'] = False
ap.box_service.available = True
session = _make_session(
mcp_module,
{
'name': 'blocked',
'uuid': 'blocked-uuid',
'mode': 'stdio',
'command': 'python',
'args': [],
'env': {},
},
ap=ap,
)
session._box_stdio_runtime.initialize = AsyncMock()
with pytest.raises(RuntimeError, match='disabled by instance policy'):
await session._init_stdio_python_server()
session._box_stdio_runtime.initialize.assert_not_awaited()
@pytest.mark.asyncio
async def test_init_box_stdio_server_stages_host_path_in_shared_workspace(mcp_module, tmp_path):
mcp_stdio_module = sys.modules['langbot.pkg.provider.tools.loaders.mcp_stdio']
@@ -931,12 +990,16 @@ async def test_init_box_stdio_server_stages_host_path_in_shared_workspace(mcp_mo
async def initialize(self):
return None
captured_transport = {}
@asynccontextmanager
async def fake_websocket_client(_url: str):
async def fake_authenticated_websocket_client(url: str, headers: dict[str, str]):
captured_transport['url'] = url
captured_transport['headers'] = headers
yield ('read-stream', 'write-stream')
mcp_stdio_module.ClientSession = FakeClientSession
mcp_stdio_module.websocket_client = fake_websocket_client
mcp_stdio_module.authenticated_websocket_client = fake_authenticated_websocket_client
ap = _make_ap()
ap.box_service.available = True
@@ -947,7 +1010,17 @@ async def test_init_box_stdio_server_stages_host_path_in_shared_workspace(mcp_mo
execute=AsyncMock(return_value=SimpleNamespace(ok=True, stderr='', exit_code=0))
)
ap.box_service.start_managed_process = AsyncMock(return_value={})
ap.box_service.get_managed_process_websocket_url = Mock(return_value='ws://box.example/process')
ap.box_service.get_managed_process_websocket_connection = AsyncMock(
return_value=(
'ws://box.example/process',
{
'X-LangBot-Box-Control-Token': 'secret-token',
'X-LangBot-Instance-Id': 'instance-a',
'X-LangBot-Workspace-Id': 'workspace-a',
'X-LangBot-Placement-Generation': '1',
},
)
)
host_path = tmp_path / 'mcp-source'
host_path.mkdir()
@@ -971,7 +1044,8 @@ async def test_init_box_stdio_server_stages_host_path_in_shared_workspace(mcp_mo
await session.exit_stack.aclose()
assert ap.box_service.create_session.await_count == 1
session_payload = ap.box_service.create_session.await_args.args[0]
assert ap.box_service.create_session.await_args.args[0] == session.execution_context
session_payload = ap.box_service.create_session.await_args.args[1]
assert session_payload['session_id'] == 'mcp-shared'
assert 'host_path' not in session_payload
assert ap.box_service.build_spec.call_count == 1
@@ -981,11 +1055,22 @@ async def test_init_box_stdio_server_stages_host_path_in_shared_workspace(mcp_mo
staged_file = tmp_path / 'shared-box-workspace' / '.mcp' / 'u1' / 'workspace' / 'server.py'
assert staged_file.read_text(encoding='utf-8') == 'print("hello")\n'
process_payload = ap.box_service.start_managed_process.await_args.args[1]
assert ap.box_service.start_managed_process.await_args.args[0] == session.execution_context
process_payload = ap.box_service.start_managed_process.await_args.args[2]
assert process_payload['process_id'] == 'u1'
assert process_payload['command'] == 'python'
assert process_payload['args'] == ['/workspace/.mcp/u1/workspace/server.py']
assert process_payload['cwd'] == '/workspace/.mcp/u1/workspace'
assert captured_transport == {
'url': 'ws://box.example/process',
'headers': {
'X-LangBot-Box-Control-Token': 'secret-token',
'X-LangBot-Instance-Id': 'instance-a',
'X-LangBot-Workspace-Id': 'workspace-a',
'X-LangBot-Placement-Generation': '1',
},
}
assert 'secret-token' not in captured_transport['url']
@pytest.mark.asyncio
@@ -1024,7 +1109,7 @@ async def test_stdio_handshake_raises_coldstart_retry_while_process_alive(mcp_mo
ap.box_service.available = True
ap.box_service.create_session = AsyncMock(return_value={})
ap.box_service.start_managed_process = AsyncMock(return_value={})
ap.box_service.get_managed_process_websocket_url = Mock(return_value='ws://box/p')
ap.box_service.get_managed_process_websocket_connection = AsyncMock(return_value=('ws://box/p', {}))
session = _make_session(
mcp_module,
@@ -1090,7 +1175,7 @@ async def test_stdio_handshake_raises_fatal_when_process_exited(mcp_module, tmp_
ap.box_service.available = True
ap.box_service.create_session = AsyncMock(return_value={})
ap.box_service.start_managed_process = AsyncMock(return_value={})
ap.box_service.get_managed_process_websocket_url = Mock(return_value='ws://box/p')
ap.box_service.get_managed_process_websocket_connection = AsyncMock(return_value=('ws://box/p', {}))
session = _make_session(
mcp_module,
tests/unit_tests/provider/test_mcp_remote_transport.py
+24 -1
View File
@@ -12,9 +12,17 @@ import pytest
from aiohttp import web
from mcp import types as mcp_types
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.provider.tools.loaders.mcp import MCPToolCallTimeoutError, RuntimeMCPSession
TEST_EXECUTION_CONTEXT = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
class _TransportProbe:
def __init__(self, streamable_status: int | None) -> None:
self.streamable_status = streamable_status
@@ -151,7 +159,21 @@ def _session(
timeout: float = 2,
tool_call_timeout_sec: float = 300,
) -> RuntimeMCPSession:
app = cast(Any, SimpleNamespace(logger=Mock()))
app = cast(
Any,
SimpleNamespace(
logger=Mock(),
workspace_service=SimpleNamespace(
get_execution_binding=AsyncMock(
return_value=SimpleNamespace(
instance_uuid=TEST_EXECUTION_CONTEXT.instance_uuid,
workspace_uuid=TEST_EXECUTION_CONTEXT.workspace_uuid,
placement_generation=TEST_EXECUTION_CONTEXT.placement_generation,
)
)
),
),
)
return RuntimeMCPSession(
'remote-transport-test',
{
@@ -163,6 +185,7 @@ def _session(
},
True,
app,
TEST_EXECUTION_CONTEXT,
)
tests/unit_tests/provider/test_mcp_resources.py
+369 -13
View File
@@ -11,6 +11,7 @@ import pytest
from mcp import types as mcp_types
from mcp.shared.exceptions import McpError
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.provider.tools.loaders.mcp import (
MCP_RESOURCE_CONTEXT_QUERY_KEY,
MCP_RESOURCE_TRACE_QUERY_KEY,
@@ -23,10 +24,30 @@ from langbot.pkg.provider.tools.loaders.mcp import (
RuntimeMCPSession,
)
from langbot.pkg.telemetry import features as telemetry_features
from langbot.pkg.workspace.errors import WorkspaceGenerationMismatchError
TEST_EXECUTION_CONTEXT = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
query_uuid='query-a',
)
def _app() -> SimpleNamespace:
return SimpleNamespace(logger=Mock())
return SimpleNamespace(
logger=Mock(),
workspace_service=SimpleNamespace(
get_execution_binding=AsyncMock(
return_value=SimpleNamespace(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
)
),
)
def _connected_session(
@@ -35,8 +56,15 @@ def _connected_session(
uuid: str = 'srv-1',
resources: list[dict] | None = None,
templates: list[dict] | None = None,
execution_context: ExecutionContext = TEST_EXECUTION_CONTEXT,
) -> RuntimeMCPSession:
session = RuntimeMCPSession(name, {'uuid': uuid, 'mode': 'remote'}, True, _app())
session = RuntimeMCPSession(
name,
{'uuid': uuid, 'mode': 'remote'},
True,
_app(),
execution_context,
)
session.status = MCPSessionStatus.CONNECTED
session.session = SimpleNamespace(read_resource=AsyncMock())
session.resources = resources or [
@@ -56,8 +84,24 @@ def _connected_session(
return session
def _query() -> SimpleNamespace:
return SimpleNamespace(variables={})
def _query(variables: dict | None = None, context: ExecutionContext = TEST_EXECUTION_CONTEXT) -> SimpleNamespace:
return SimpleNamespace(
instance_uuid=context.instance_uuid,
workspace_uuid=context.workspace_uuid,
placement_generation=context.placement_generation,
bot_uuid=context.bot_uuid,
pipeline_uuid=context.pipeline_uuid,
query_uuid=context.query_uuid,
variables=variables or {},
)
def _register_session(loader: MCPLoader, session: RuntimeMCPSession) -> None:
loader._register_session(
session.execution_context,
session.server_name,
session,
)
def _http_status_error(status_code: int) -> httpx.HTTPStatusError:
@@ -84,6 +128,7 @@ async def test_invoke_mcp_tool_uses_configurable_request_timeout():
},
True,
_app(),
TEST_EXECUTION_CONTEXT,
)
session.session = SimpleNamespace(call_tool=AsyncMock(return_value=_tool_result()))
@@ -108,6 +153,7 @@ async def test_invoke_mcp_tool_zero_timeout_disables_request_deadline():
},
True,
_app(),
TEST_EXECUTION_CONTEXT,
)
session.session = SimpleNamespace(call_tool=AsyncMock(return_value=_tool_result()))
@@ -131,6 +177,7 @@ async def test_invoke_mcp_tool_timeout_is_not_retried_and_session_remains_usable
},
True,
_app(),
TEST_EXECUTION_CONTEXT,
)
timeout = McpError(
mcp_types.ErrorData(
@@ -165,6 +212,7 @@ def test_invalid_tool_call_timeout_falls_back_to_default(invalid_timeout):
},
True,
ap,
TEST_EXECUTION_CONTEXT,
)
assert session.tool_call_timeout_sec == MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS
@@ -178,6 +226,7 @@ async def test_remote_transport_falls_back_to_sse_for_compatible_http_status_in_
{'uuid': 'srv-1', 'mode': 'remote', 'url': 'https://example.com/mcp'},
True,
_app(),
TEST_EXECUTION_CONTEXT,
)
session._init_streamable_http_server = AsyncMock(
side_effect=ExceptionGroup('transport failed', [_http_status_error(405)])
@@ -197,6 +246,7 @@ async def test_remote_transport_does_not_fallback_for_auth_http_status():
{'uuid': 'srv-1', 'mode': 'remote', 'url': 'https://example.com/mcp'},
True,
_app(),
TEST_EXECUTION_CONTEXT,
)
error = _http_status_error(403)
session._init_streamable_http_server = AsyncMock(side_effect=error)
@@ -354,10 +404,18 @@ def test_resource_uri_allowed_supports_listed_templates_conservatively():
async def test_mcp_loader_can_hide_synthetic_resource_tools():
loader = MCPLoader(_app())
session = _connected_session()
loader.sessions = {'docs': session}
_register_session(loader, session)
with_resource_tools = await loader.get_tools(['srv-1'], include_resource_tools=True)
without_resource_tools = await loader.get_tools(['srv-1'], include_resource_tools=False)
with_resource_tools = await loader.get_tools(
TEST_EXECUTION_CONTEXT,
['srv-1'],
include_resource_tools=True,
)
without_resource_tools = await loader.get_tools(
TEST_EXECUTION_CONTEXT,
['srv-1'],
include_resource_tools=False,
)
assert {tool.name for tool in with_resource_tools} == {
MCP_TOOL_LIST_RESOURCES,
@@ -370,9 +428,9 @@ async def test_mcp_loader_can_hide_synthetic_resource_tools():
async def test_mcp_loader_refuses_resource_tool_calls_when_agent_read_disabled():
loader = MCPLoader(_app())
session = _connected_session()
loader.sessions = {'docs': session}
query = SimpleNamespace(
variables={
_register_session(loader, session)
query = _query(
{
'_pipeline_bound_mcp_servers': ['srv-1'],
'_pipeline_mcp_resource_agent_read_enabled': False,
}
@@ -411,9 +469,10 @@ async def test_build_resource_context_for_query_uses_only_bound_attached_text_re
)
]
)
loader.sessions = {'docs': docs, 'other': other}
query = SimpleNamespace(
variables={
_register_session(loader, docs)
_register_session(loader, other)
query = _query(
{
'_pipeline_bound_mcp_servers': ['srv-1'],
'_pipeline_mcp_resource_attachments': [
{'server_uuid': 'srv-1', 'server_name': 'docs', 'uri': 'file:///README.md', 'mode': 'pinned'},
@@ -433,6 +492,140 @@ async def test_build_resource_context_for_query_uses_only_bound_attached_text_re
other.session.read_resource.assert_not_called()
def test_mcp_loader_session_keys_do_not_collide_between_workspaces():
loader = MCPLoader(_app())
workspace_b = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-b',
placement_generation=1,
query_uuid='query-b',
)
session_a = _connected_session(name='docs', uuid='srv-a')
session_b = _connected_session(
name='docs',
uuid='srv-b',
execution_context=workspace_b,
)
_register_session(loader, session_a)
_register_session(loader, session_b)
assert len(loader.sessions) == 2
assert loader.get_session(TEST_EXECUTION_CONTEXT, 'docs') is session_a
assert loader.get_session(workspace_b, 'docs') is session_b
assert loader.get_session(TEST_EXECUTION_CONTEXT, 'docs') is not loader.get_session(workspace_b, 'docs')
@pytest.mark.asyncio
async def test_mcp_tool_result_is_discarded_when_generation_changes_during_call():
session = _connected_session()
binding = SimpleNamespace(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
session.ap.workspace_service.get_execution_binding.side_effect = [
binding,
binding,
WorkspaceGenerationMismatchError('generation changed during tool call'),
]
session.session = SimpleNamespace(call_tool=AsyncMock(return_value=SimpleNamespace(isError=False, content=[])))
with pytest.raises(WorkspaceGenerationMismatchError):
await session.invoke_mcp_tool('side_effecting_tool', {})
session.session.call_tool.assert_awaited_once_with(
'side_effecting_tool',
{},
read_timeout_seconds=timedelta(seconds=MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS),
)
@pytest.mark.asyncio
async def test_mcp_resource_cache_is_not_served_to_stale_generation():
session = _connected_session()
session._resource_cache[('file:///README.md', 10, None, False)] = {
'cached_at': 0,
'envelope': {'contents': [{'type': 'text', 'text': 'stale'}]},
}
session.ap.workspace_service.get_execution_binding.side_effect = WorkspaceGenerationMismatchError(
'stale generation'
)
with pytest.raises(WorkspaceGenerationMismatchError):
await session.read_resource_envelope('file:///README.md', max_bytes=10)
session.session.read_resource.assert_not_awaited()
@pytest.mark.asyncio
async def test_directory_projection_retires_idle_mcp_scope_without_db_poll():
loader = MCPLoader(_app())
sessions = []
for index in range(100):
context = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid=f'workspace-{index}',
placement_generation=1,
)
session = RuntimeMCPSession(
f'server-{index}',
{'uuid': f'srv-{index}', 'mode': 'remote'},
True,
loader.ap,
context,
)
session.shutdown = AsyncMock()
loader._register_session(context, session.server_name, session)
sessions.append(session)
loader.reconcile_execution_projection('instance-a', {})
reconcile_task = loader._projection_reconcile_task
assert reconcile_task is not None
assert len(loader._pending_projection_retirements) == 100
# A second projection coalesces into the same worker instead of creating
# one timer or task per Workspace.
loader.reconcile_execution_projection('instance-a', {})
assert loader._projection_reconcile_task is reconcile_task
await asyncio.wait_for(reconcile_task, timeout=1)
assert loader.sessions == {}
assert loader._scope_generations == {}
assert loader._pending_projection_retirements == set()
assert sum(session.shutdown.await_count for session in sessions) == 100
loader.ap.workspace_service.get_execution_binding.assert_not_awaited()
@pytest.mark.asyncio
async def test_directory_projection_keeps_matching_and_unaffected_mcp_scopes():
loader = MCPLoader(_app())
matching = _connected_session()
other_context = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-b',
placement_generation=1,
)
unaffected = _connected_session(
name='other',
uuid='srv-2',
execution_context=other_context,
)
_register_session(loader, matching)
_register_session(loader, unaffected)
loader.reconcile_execution_projection(
'instance-a',
{'workspace-a': 1},
affected_workspace_uuids={'workspace-a'},
)
await asyncio.sleep(0)
assert loader.get_session(TEST_EXECUTION_CONTEXT, 'docs') is matching
assert loader.get_session(other_context, 'other') is unaffected
assert loader._projection_reconcile_task is None
@pytest.mark.asyncio
async def test_mcp_loader_shutdown_cancels_startup_tasks_and_closes_sessions_concurrently():
loader = MCPLoader(_app())
@@ -454,6 +647,7 @@ async def test_mcp_loader_shutdown_cancels_startup_tasks_and_closes_sessions_con
class Session:
def __init__(self, name: str):
self.name = name
self.server_name = name
async def shutdown(self):
started.add(self.name)
@@ -470,3 +664,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
tests/unit_tests/provider/test_mcp_stdio_policy.py
+70
View File
@@ -0,0 +1,70 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
from langbot.pkg.provider.tools.loaders.mcp_policy import (
MCPStdioDisabledError,
require_stdio_mcp_enabled,
stdio_mcp_enabled,
)
from langbot.pkg.provider.tools.loaders.mcp import MCPLoader
def _app(config: dict) -> SimpleNamespace:
return SimpleNamespace(instance_config=SimpleNamespace(data=config))
def test_oss_default_remains_enabled_when_key_is_absent():
assert stdio_mcp_enabled(_app({})) is True
@pytest.mark.parametrize(
'value',
[False, 'false', 0, None, {}, []],
)
def test_disabled_or_invalid_values_fail_closed(value):
ap = _app({'mcp': {'stdio': {'enabled': value}}})
assert stdio_mcp_enabled(ap) is False
with pytest.raises(MCPStdioDisabledError, match='disabled by instance policy'):
require_stdio_mcp_enabled(ap, {'mode': 'stdio'})
def test_remote_transport_is_independent_of_stdio_gate():
ap = _app({'mcp': {'stdio': {'enabled': False}}})
require_stdio_mcp_enabled(ap, {'mode': 'remote'})
@pytest.mark.asyncio
async def test_bootstrap_retains_but_does_not_launch_disabled_stdio_rows():
server = SimpleNamespace(uuid='server-a', workspace_uuid='workspace-a')
result = Mock()
result.all.return_value = [server]
ap = _app({'mcp': {'stdio': {'enabled': False}}})
ap.logger = Mock()
ap.persistence_mgr = SimpleNamespace(
execute_async=AsyncMock(return_value=result),
serialize_model=Mock(
return_value={
'uuid': 'server-a',
'workspace_uuid': 'workspace-a',
'name': 'local',
'mode': 'stdio',
'enable': True,
'extra_args': {},
}
),
)
ap.workspace_service = SimpleNamespace(get_execution_binding=AsyncMock())
loader = MCPLoader(ap)
loader.host_mcp_server = AsyncMock()
await loader.load_mcp_servers_from_db()
loader.host_mcp_server.assert_not_awaited()
ap.workspace_service.get_execution_binding.assert_not_awaited()
assert loader.sessions == {}
tests/unit_tests/provider/test_model_manager.py
+317 -45
View File
@@ -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
@@ -16,7 +17,15 @@ from langbot.pkg.provider.modelmgr import requester
from langbot.pkg.entity.persistence import model as persistence_model
from langbot.pkg.entity.errors import provider as provider_errors
from langbot.pkg.provider.modelmgr import token
from tests.unit_tests.provider.conftest import _make_mock_result, _make_row_mock
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.workspace.entities import WorkspaceExecutionBinding
from langbot.pkg.workspace.errors import WorkspaceGenerationMismatchError, WorkspaceInvariantError
from tests.unit_tests.provider.conftest import (
TEST_EXECUTION_CONTEXT,
TEST_WORKSPACE_UUID,
_make_mock_result,
_make_row_mock,
)
# ============================================================================
@@ -63,6 +72,20 @@ async def test_model_manager_skips_space_sync_when_disabled(mock_app_for_modelmg
app.space_service.get_models.assert_not_called()
@pytest.mark.asyncio
async def test_model_manager_skips_legacy_space_sync_in_cloud_runtime(mock_app_for_modelmgr):
"""Cloud startup must not resolve an OSS-local Workspace for legacy model sync."""
app = mock_app_for_modelmgr
app.instance_config.data = {'space': {'disable_models_service': False}}
app.persistence_mgr.mode = SimpleNamespace(value='cloud_runtime')
model_mgr = ModelManager(app)
model_mgr.load_models_from_db = AsyncMock()
await model_mgr.initialize()
app.workspace_service.get_local_execution_binding.assert_not_awaited()
@pytest.mark.asyncio
async def test_sync_new_models_from_space_creates_rerank_models(mock_app_for_modelmgr):
"""Space rerank entries are discovered and persisted under the shared provider."""
@@ -91,9 +114,10 @@ async def test_sync_new_models_from_space_creates_rerank_models(mock_app_for_mod
app.rerank_models_service.get_rerank_models = AsyncMock(return_value=[])
model_mgr = ModelManager(app)
await model_mgr.sync_new_models_from_space()
await model_mgr.sync_new_models_from_space(TEST_EXECUTION_CONTEXT)
app.rerank_models_service.create_rerank_model.assert_awaited_once_with(
TEST_EXECUTION_CONTEXT,
{
'uuid': 'rerank-model-uuid',
'name': 'Qwen3-Reranker-8B',
@@ -134,13 +158,33 @@ async def test_model_manager_load_models_from_db(fake_requester_registry, fake_p
# Check providers loaded
assert len(model_mgr.provider_dict) == 2
assert fake_persistence_data['provider_uuid'] in model_mgr.provider_dict
assert fake_persistence_data['provider_uuid2'] in model_mgr.provider_dict
assert {provider.provider_entity.uuid for provider in model_mgr.provider_dict.values()} == {
fake_persistence_data['provider_uuid'],
fake_persistence_data['provider_uuid2'],
}
# Check models loaded
assert len(model_mgr.llm_models) == 2
assert len(model_mgr.embedding_models) == 1
assert len(model_mgr.rerank_models) == 1
assert len(model_mgr.llm_model_dict) == 2
assert len(model_mgr.embedding_model_dict) == 1
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
@@ -161,7 +205,7 @@ async def test_model_manager_load_provider_unknown_requester(mock_app_for_modelm
}
with pytest.raises(provider_errors.RequesterNotFoundError) as exc_info:
await model_mgr.load_provider(provider_info)
await model_mgr.load_provider(TEST_EXECUTION_CONTEXT, provider_info)
assert exc_info.value.requester_name == 'non-existent-requester'
@@ -180,7 +224,7 @@ async def test_model_manager_load_provider_from_dict(fake_requester_registry):
'api_keys': ['dict-key'],
}
runtime_provider = await model_mgr.load_provider(provider_info)
runtime_provider = await model_mgr.load_provider(TEST_EXECUTION_CONTEXT, provider_info)
assert runtime_provider.provider_entity.uuid == 'dict-provider-uuid'
assert runtime_provider.provider_entity.name == 'Dict Provider'
@@ -197,7 +241,7 @@ async def test_model_manager_load_provider_from_entity(fake_requester_registry,
provider_entity = fake_persistence_data['providers'][0]
runtime_provider = await model_mgr.load_provider(provider_entity)
runtime_provider = await model_mgr.load_provider(TEST_EXECUTION_CONTEXT, provider_entity)
assert runtime_provider.provider_entity.uuid == provider_entity.uuid
assert runtime_provider.requester is not None
@@ -224,7 +268,7 @@ async def test_model_manager_get_model_by_uuid(fake_requester_registry, fake_per
model_mgr.ap.persistence_mgr.execute_async = fake_execute
await model_mgr.initialize()
model = await model_mgr.get_model_by_uuid('test-llm-uuid-1')
model = await model_mgr.get_model_by_uuid(TEST_EXECUTION_CONTEXT, 'test-llm-uuid-1')
assert model.model_entity.uuid == 'test-llm-uuid-1'
assert model.model_entity.name == 'TestLLM-1'
@@ -237,7 +281,7 @@ async def test_model_manager_get_model_by_uuid_not_found(fake_requester_registry
await model_mgr.initialize()
with pytest.raises(ValueError) as exc_info:
await model_mgr.get_model_by_uuid('unknown-model-uuid')
await model_mgr.get_model_by_uuid(TEST_EXECUTION_CONTEXT, 'unknown-model-uuid')
assert 'unknown-model-uuid' in str(exc_info.value)
@@ -258,7 +302,10 @@ async def test_model_manager_get_embedding_model_by_uuid(fake_requester_registry
model_mgr.ap.persistence_mgr.execute_async = fake_execute
await model_mgr.initialize()
model = await model_mgr.get_embedding_model_by_uuid('test-embedding-uuid-1')
model = await model_mgr.get_embedding_model_by_uuid(
TEST_EXECUTION_CONTEXT,
'test-embedding-uuid-1',
)
assert model.model_entity.uuid == 'test-embedding-uuid-1'
@@ -270,7 +317,10 @@ async def test_model_manager_get_embedding_model_by_uuid_not_found(fake_requeste
await model_mgr.initialize()
with pytest.raises(ValueError):
await model_mgr.get_embedding_model_by_uuid('unknown-embedding-uuid')
await model_mgr.get_embedding_model_by_uuid(
TEST_EXECUTION_CONTEXT,
'unknown-embedding-uuid',
)
@pytest.mark.asyncio
@@ -289,7 +339,7 @@ async def test_model_manager_get_rerank_model_by_uuid(fake_requester_registry, f
model_mgr.ap.persistence_mgr.execute_async = fake_execute
await model_mgr.initialize()
model = await model_mgr.get_rerank_model_by_uuid('test-rerank-uuid-1')
model = await model_mgr.get_rerank_model_by_uuid(TEST_EXECUTION_CONTEXT, 'test-rerank-uuid-1')
assert model.model_entity.uuid == 'test-rerank-uuid-1'
@@ -301,7 +351,7 @@ async def test_model_manager_get_rerank_model_by_uuid_not_found(fake_requester_r
await model_mgr.initialize()
with pytest.raises(ValueError):
await model_mgr.get_rerank_model_by_uuid('unknown-rerank-uuid')
await model_mgr.get_rerank_model_by_uuid(TEST_EXECUTION_CONTEXT, 'unknown-rerank-uuid')
# ============================================================================
@@ -325,12 +375,12 @@ async def test_model_manager_remove_llm_model(fake_requester_registry, fake_pers
model_mgr.ap.persistence_mgr.execute_async = fake_execute
await model_mgr.initialize()
assert len(model_mgr.llm_models) == 2
assert len(model_mgr.llm_model_dict) == 2
await model_mgr.remove_llm_model('test-llm-uuid-1')
await model_mgr.remove_llm_model(TEST_EXECUTION_CONTEXT, 'test-llm-uuid-1')
assert len(model_mgr.llm_models) == 1
assert model_mgr.llm_models[0].model_entity.uuid == 'test-llm-uuid-2'
assert len(model_mgr.llm_model_dict) == 1
assert next(iter(model_mgr.llm_model_dict.values())).model_entity.uuid == 'test-llm-uuid-2'
@pytest.mark.asyncio
@@ -349,12 +399,12 @@ async def test_model_manager_remove_llm_model_not_found(fake_requester_registry,
model_mgr.ap.persistence_mgr.execute_async = fake_execute
await model_mgr.initialize()
original_count = len(model_mgr.llm_models)
original_count = len(model_mgr.llm_model_dict)
# Removing unknown model should do nothing (no error)
await model_mgr.remove_llm_model('unknown-model-uuid')
await model_mgr.remove_llm_model(TEST_EXECUTION_CONTEXT, 'unknown-model-uuid')
assert len(model_mgr.llm_models) == original_count
assert len(model_mgr.llm_model_dict) == original_count
@pytest.mark.asyncio
@@ -373,11 +423,11 @@ async def test_model_manager_remove_embedding_model(fake_requester_registry, fak
model_mgr.ap.persistence_mgr.execute_async = fake_execute
await model_mgr.initialize()
assert len(model_mgr.embedding_models) == 1
assert len(model_mgr.embedding_model_dict) == 1
await model_mgr.remove_embedding_model('test-embedding-uuid-1')
await model_mgr.remove_embedding_model(TEST_EXECUTION_CONTEXT, 'test-embedding-uuid-1')
assert len(model_mgr.embedding_models) == 0
assert len(model_mgr.embedding_model_dict) == 0
@pytest.mark.asyncio
@@ -396,11 +446,11 @@ async def test_model_manager_remove_rerank_model(fake_requester_registry, fake_p
model_mgr.ap.persistence_mgr.execute_async = fake_execute
await model_mgr.initialize()
assert len(model_mgr.rerank_models) == 1
assert len(model_mgr.rerank_model_dict) == 1
await model_mgr.remove_rerank_model('test-rerank-uuid-1')
await model_mgr.remove_rerank_model(TEST_EXECUTION_CONTEXT, 'test-rerank-uuid-1')
assert len(model_mgr.rerank_models) == 0
assert len(model_mgr.rerank_model_dict) == 0
@pytest.mark.asyncio
@@ -419,11 +469,17 @@ async def test_model_manager_remove_provider(fake_requester_registry, fake_persi
model_mgr.ap.persistence_mgr.execute_async = fake_execute
await model_mgr.initialize()
assert fake_persistence_data['provider_uuid'] in model_mgr.provider_dict
assert any(
provider.provider_entity.uuid == fake_persistence_data['provider_uuid']
for provider in model_mgr.provider_dict.values()
)
await model_mgr.remove_provider(fake_persistence_data['provider_uuid'])
await model_mgr.remove_provider(TEST_EXECUTION_CONTEXT, fake_persistence_data['provider_uuid'])
assert fake_persistence_data['provider_uuid'] not in model_mgr.provider_dict
assert all(
provider.provider_entity.uuid != fake_persistence_data['provider_uuid']
for provider in model_mgr.provider_dict.values()
)
# ============================================================================
@@ -541,7 +597,7 @@ async def test_model_manager_init_temporary_runtime_llm_model(fake_requester_reg
'extra_args': {'temperature': 0.5},
}
runtime_model = await model_mgr.init_temporary_runtime_llm_model(model_info)
runtime_model = await model_mgr.init_temporary_runtime_llm_model(TEST_EXECUTION_CONTEXT, model_info)
assert runtime_model.model_entity.uuid == 'temp-model-uuid'
assert runtime_model.model_entity.name == 'TempModel'
@@ -571,7 +627,10 @@ async def test_model_manager_init_temporary_runtime_embedding_model(fake_request
'extra_args': {'dimensions': 512},
}
runtime_model = await model_mgr.init_temporary_runtime_embedding_model(model_info)
runtime_model = await model_mgr.init_temporary_runtime_embedding_model(
TEST_EXECUTION_CONTEXT,
model_info,
)
assert runtime_model.model_entity.uuid == 'temp-embedding-uuid'
assert runtime_model.model_entity.name == 'TempEmbedding'
@@ -596,7 +655,10 @@ async def test_model_manager_init_temporary_runtime_rerank_model(fake_requester_
'extra_args': {},
}
runtime_model = await model_mgr.init_temporary_runtime_rerank_model(model_info)
runtime_model = await model_mgr.init_temporary_runtime_rerank_model(
TEST_EXECUTION_CONTEXT,
model_info,
)
assert runtime_model.model_entity.uuid == 'temp-rerank-uuid'
assert runtime_model.model_entity.name == 'TempRerank'
@@ -632,12 +694,16 @@ async def test_model_manager_reload_provider(fake_requester_registry, fake_persi
model_mgr.ap.persistence_mgr.execute_async = fake_execute
await model_mgr.initialize()
original_provider = model_mgr.provider_dict[fake_persistence_data['provider_uuid']]
original_provider = await model_mgr.get_provider_by_uuid(
TEST_EXECUTION_CONTEXT,
fake_persistence_data['provider_uuid'],
)
original_base_url = original_provider.provider_entity.base_url
# Setup for reload - return updated provider
async def reload_execute(query):
updated_provider = persistence_model.ModelProvider(
workspace_uuid=TEST_WORKSPACE_UUID,
uuid=fake_persistence_data['provider_uuid'],
name='Updated Provider',
requester='fake-requester',
@@ -648,9 +714,12 @@ async def test_model_manager_reload_provider(fake_requester_registry, fake_persi
model_mgr.ap.persistence_mgr.execute_async = reload_execute
await model_mgr.reload_provider(fake_persistence_data['provider_uuid'])
await model_mgr.reload_provider(TEST_EXECUTION_CONTEXT, fake_persistence_data['provider_uuid'])
updated_provider = model_mgr.provider_dict[fake_persistence_data['provider_uuid']]
updated_provider = await model_mgr.get_provider_by_uuid(
TEST_EXECUTION_CONTEXT,
fake_persistence_data['provider_uuid'],
)
assert updated_provider.provider_entity.base_url == 'https://updated.example.com'
assert updated_provider.provider_entity.base_url != original_base_url
@@ -667,7 +736,7 @@ async def test_model_manager_reload_provider_not_found(fake_requester_registry):
model_mgr.ap.persistence_mgr.execute_async = fake_execute
with pytest.raises(provider_errors.ProviderNotFoundError) as exc_info:
await model_mgr.reload_provider('unknown-provider-uuid')
await model_mgr.reload_provider(TEST_EXECUTION_CONTEXT, 'unknown-provider-uuid')
assert exc_info.value.provider_name == 'unknown-provider-uuid'
@@ -686,7 +755,11 @@ async def test_model_manager_load_llm_model_with_provider(
model_entity = fake_persistence_data['llm_models'][0]
runtime_model = await model_mgr.load_llm_model_with_provider(model_entity, runtime_provider)
runtime_model = await model_mgr.load_llm_model_with_provider(
TEST_EXECUTION_CONTEXT,
model_entity,
runtime_provider,
)
assert runtime_model.model_entity.uuid == model_entity.uuid
assert runtime_model.provider is runtime_provider
@@ -702,7 +775,11 @@ async def test_model_manager_load_llm_model_with_provider_from_row(
model_entity = fake_persistence_data['llm_models'][0]
row_mock = _make_row_mock(model_entity)
runtime_model = await model_mgr.load_llm_model_with_provider(row_mock, runtime_provider)
runtime_model = await model_mgr.load_llm_model_with_provider(
TEST_EXECUTION_CONTEXT,
row_mock,
runtime_provider,
)
assert runtime_model.model_entity.uuid == model_entity.uuid
@@ -716,7 +793,11 @@ async def test_model_manager_load_embedding_model_with_provider(
model_entity = fake_persistence_data['embedding_models'][0]
runtime_model = await model_mgr.load_embedding_model_with_provider(model_entity, runtime_provider)
runtime_model = await model_mgr.load_embedding_model_with_provider(
TEST_EXECUTION_CONTEXT,
model_entity,
runtime_provider,
)
assert runtime_model.model_entity.uuid == model_entity.uuid
assert runtime_model.provider is runtime_provider
@@ -735,6 +816,7 @@ async def test_model_manager_load_rerank_model_with_provider(fake_requester_regi
)
await requester_inst.initialize()
provider = requester.RuntimeProvider(
execution_context=TEST_EXECUTION_CONTEXT,
provider_entity=provider_entity,
token_mgr=token_mgr,
requester=requester_inst,
@@ -742,7 +824,11 @@ async def test_model_manager_load_rerank_model_with_provider(fake_requester_regi
model_entity = fake_persistence_data['rerank_models'][0]
runtime_model = await model_mgr.load_rerank_model_with_provider(model_entity, provider)
runtime_model = await model_mgr.load_rerank_model_with_provider(
TEST_EXECUTION_CONTEXT,
model_entity,
provider,
)
assert runtime_model.model_entity.uuid == model_entity.uuid
assert runtime_model.provider is provider
@@ -766,6 +852,7 @@ async def test_model_manager_logs_warning_for_missing_provider(fake_requester_re
elif 'llm_models' in query_str:
# Return model with missing provider
fake_model = persistence_model.LLMModel(
workspace_uuid=TEST_WORKSPACE_UUID,
uuid='model-with-missing-provider',
name='MissingProviderModel',
provider_uuid='missing-provider-uuid',
@@ -779,7 +866,7 @@ async def test_model_manager_logs_warning_for_missing_provider(fake_requester_re
await model_mgr.initialize()
# Should have logged warning and skipped the model
assert len(model_mgr.llm_models) == 0
assert len(model_mgr.llm_model_dict) == 0
model_mgr.ap.logger.warning.assert_called()
@@ -793,6 +880,7 @@ async def test_model_manager_handles_requester_not_found_gracefully(fake_request
if 'model_providers' in query_str:
# Return provider with unknown requester
fake_provider = persistence_model.ModelProvider(
workspace_uuid=TEST_WORKSPACE_UUID,
uuid='provider-with-unknown-requester',
name='Unknown Requester Provider',
requester='unknown-requester-name',
@@ -802,6 +890,7 @@ async def test_model_manager_handles_requester_not_found_gracefully(fake_request
return _make_mock_result([_make_row_mock(fake_provider)])
elif 'llm_models' in query_str:
fake_model = persistence_model.LLMModel(
workspace_uuid=TEST_WORKSPACE_UUID,
uuid='model-uuid',
name='Model',
provider_uuid='provider-with-unknown-requester',
@@ -816,7 +905,7 @@ async def test_model_manager_handles_requester_not_found_gracefully(fake_request
# Provider should be skipped
assert len(model_mgr.provider_dict) == 0
assert len(model_mgr.llm_models) == 0
assert len(model_mgr.llm_model_dict) == 0
model_mgr.ap.logger.warning.assert_called()
@@ -833,6 +922,189 @@ def test_requester_not_found_error_str():
assert error.requester_name == 'test-requester'
@pytest.mark.asyncio
async def test_runtime_cache_isolates_same_resource_uuid_between_workspaces(fake_requester_registry):
"""A UUID collision cannot select another Workspace's runtime object."""
model_mgr = fake_requester_registry
await model_mgr.initialize()
contexts = {
workspace_uuid: ExecutionContext(
instance_uuid='test-instance',
workspace_uuid=workspace_uuid,
placement_generation=1,
)
for workspace_uuid in ('workspace-a', 'workspace-b')
}
async def resolve_binding(workspace_uuid, *, expected_generation=None):
assert expected_generation in (None, 1)
return WorkspaceExecutionBinding(
instance_uuid='test-instance',
workspace_uuid=workspace_uuid,
placement_generation=1,
write_fenced=False,
state='active',
)
model_mgr.ap.workspace_service.get_execution_binding = AsyncMock(side_effect=resolve_binding)
for workspace_uuid, context in contexts.items():
provider = await model_mgr.load_provider(
context,
{
'uuid': 'shared-provider',
'name': f'Provider {workspace_uuid}',
'requester': 'fake-requester',
'base_url': f'https://{workspace_uuid}.example.com',
'api_keys': [],
},
)
await model_mgr.cache_provider(context, provider)
runtime_model = await model_mgr.load_llm_model_with_provider(
context,
persistence_model.LLMModel(
workspace_uuid=workspace_uuid,
uuid='shared-model',
name=f'Model {workspace_uuid}',
provider_uuid='shared-provider',
abilities=[],
extra_args={},
),
provider,
)
await model_mgr.cache_llm_model(context, runtime_model)
workspace_a_model = await model_mgr.get_model_by_uuid(contexts['workspace-a'], 'shared-model')
workspace_b_model = await model_mgr.get_model_by_uuid(contexts['workspace-b'], 'shared-model')
assert workspace_a_model.model_entity.name == 'Model workspace-a'
assert workspace_b_model.model_entity.name == 'Model workspace-b'
assert workspace_a_model is not workspace_b_model
@pytest.mark.asyncio
async def test_runtime_cache_rejects_stale_placement_generation(fake_requester_registry):
"""A stale generation is fenced before any cached model can be returned."""
model_mgr = fake_requester_registry
await model_mgr.initialize()
stale_context = TEST_EXECUTION_CONTEXT
async def reject_stale(_workspace_uuid, *, expected_generation=None):
if expected_generation == stale_context.placement_generation:
raise WorkspaceGenerationMismatchError('stale generation')
raise AssertionError('lookup must include the supplied generation')
model_mgr.ap.workspace_service.get_execution_binding = AsyncMock(side_effect=reject_stale)
with pytest.raises(WorkspaceGenerationMismatchError, match='stale generation'):
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')
tests/unit_tests/provider/test_model_service.py
+86 -18
View File
@@ -19,6 +19,8 @@ from langbot.pkg.provider.modelmgr import requester
from langbot.pkg.provider.modelmgr.modelmgr import ModelManager
from langbot.pkg.provider.modelmgr.token import TokenManager
from langbot.pkg.provider.runners.localagent import LocalAgentRunner
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.workspace.entities import WorkspaceExecutionBinding
def test_runtime_llm_model_data_preserves_uuid_after_update_payload_uuid_removed():
@@ -95,11 +97,22 @@ async def test_model_manager_initialize_skips_space_sync_after_timeout():
ap.discover = SimpleNamespace(get_components_by_kind=Mock(return_value=[]))
ap.instance_config = SimpleNamespace(data={'space': {'models_sync_timeout': 0.01}})
ap.logger = Mock()
binding = WorkspaceExecutionBinding(
instance_uuid='instance-test',
workspace_uuid='workspace-test',
placement_generation=1,
write_fenced=False,
state='active',
)
ap.workspace_service = SimpleNamespace(
get_local_execution_binding=AsyncMock(return_value=binding),
get_execution_binding=AsyncMock(return_value=binding),
)
mgr = ModelManager(ap)
mgr.load_models_from_db = AsyncMock()
async def slow_sync():
async def slow_sync(_context):
await asyncio.sleep(1)
mgr.sync_new_models_from_space = AsyncMock(side_effect=slow_sync)
@@ -117,6 +130,14 @@ async def test_updated_llm_model_is_immediately_usable_by_local_agent_pipeline()
model_uuid = 'qwen-model-uuid'
provider_uuid = 'ollama-provider-uuid'
workspace_uuid = 'workspace-test'
execution_context = ExecutionContext(
instance_uuid='instance-test',
workspace_uuid=workspace_uuid,
placement_generation=1,
bot_uuid='bot-uuid',
pipeline_uuid='pipeline-uuid',
)
ap = SimpleNamespace()
ap.logger = Mock()
@@ -126,24 +147,63 @@ async def test_updated_llm_model_is_immediately_usable_by_local_agent_pipeline()
ap.plugin_connector = SimpleNamespace(
emit_event=AsyncMock(return_value=SimpleNamespace(event=SimpleNamespace(default_prompt=[], prompt=[])))
)
binding = WorkspaceExecutionBinding(
instance_uuid='instance-test',
workspace_uuid=workspace_uuid,
placement_generation=1,
write_fenced=False,
state='active',
)
ap.workspace_service = SimpleNamespace(get_execution_binding=AsyncMock(return_value=binding))
ap.model_mgr = ModelManager(ap)
runtime_provider = Mock()
ap.model_mgr.provider_dict = {provider_uuid: runtime_provider}
ap.model_mgr.llm_models = [
requester.RuntimeLLMModel(
model_entity=persistence_model.LLMModel(
uuid=model_uuid,
name='old-qwen-name',
provider_uuid=provider_uuid,
abilities=[],
extra_args={},
),
provider=runtime_provider,
)
]
runtime_provider = Mock(
execution_context=execution_context,
provider_entity=persistence_model.ModelProvider(
workspace_uuid=workspace_uuid,
uuid=provider_uuid,
name='Ollama',
requester='ollama',
base_url='http://localhost:11434',
api_keys=[],
),
)
cache_key = ('instance-test', workspace_uuid, 1, provider_uuid)
ap.model_mgr.provider_dict = {cache_key: runtime_provider}
runtime_model = requester.RuntimeLLMModel(
execution_context=execution_context,
model_entity=persistence_model.LLMModel(
workspace_uuid=workspace_uuid,
uuid=model_uuid,
name='old-qwen-name',
provider_uuid=provider_uuid,
abilities=[],
extra_args={},
),
provider=runtime_provider,
)
ap.model_mgr.llm_model_dict = {
('instance-test', workspace_uuid, 1, model_uuid): runtime_model,
}
await LLMModelsService(ap).update_llm_model(
ap.provider_service = SimpleNamespace(
get_provider=AsyncMock(return_value={'uuid': provider_uuid, 'workspace_uuid': workspace_uuid})
)
model_service = LLMModelsService(ap)
model_service.get_llm_model = AsyncMock(
return_value={
'uuid': model_uuid,
'workspace_uuid': workspace_uuid,
'name': 'old-qwen-name',
'provider_uuid': provider_uuid,
'abilities': [],
'context_length': None,
'extra_args': {},
'prefered_ranking': 0,
}
)
await model_service.update_llm_model(
workspace_uuid,
model_uuid,
{
'name': 'Qwen3.5-27B',
@@ -153,13 +213,17 @@ async def test_updated_llm_model_is_immediately_usable_by_local_agent_pipeline()
},
)
runtime_model = await ap.model_mgr.get_model_by_uuid(model_uuid)
runtime_model = await ap.model_mgr.get_model_by_uuid(execution_context, model_uuid)
assert runtime_model.model_entity.uuid == model_uuid
assert runtime_model.model_entity.name == 'Qwen3.5-27B'
session = SimpleNamespace(
session = provider_session.Session(
instance_uuid='instance-test',
workspace_uuid=workspace_uuid,
placement_generation=1,
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
bot_uuid='bot-uuid',
)
conversation = SimpleNamespace(
uuid='conversation-uuid',
@@ -194,6 +258,9 @@ async def test_updated_llm_model_is_immediately_usable_by_local_agent_pipeline()
'output': {'misc': {'remove-think': False}},
}
query = pipeline_query.Query.model_construct(
instance_uuid='instance-test',
workspace_uuid=workspace_uuid,
placement_generation=1,
query_id='query-id',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
@@ -215,6 +282,7 @@ async def test_updated_llm_model_is_immediately_usable_by_local_agent_pipeline()
resp_message_chain=None,
current_stage_name=None,
)
object.__setattr__(query, '_execution_context', execution_context)
result = await PreProcessor(ap).process(query, 'PreProcessor')
processed_query = result.new_query
tests/unit_tests/provider/test_requester_base.py
+27 -3
View File
@@ -15,6 +15,7 @@ from langbot.pkg.provider.modelmgr import requester
from langbot.pkg.provider.modelmgr import token
from langbot.pkg.entity.persistence import model as persistence_model
from langbot.pkg.provider.modelmgr.errors import RequesterError
from tests.unit_tests.provider.conftest import TEST_EXECUTION_CONTEXT, TEST_WORKSPACE_UUID
# ============================================================================
@@ -134,6 +135,7 @@ async def test_requester_invoke_rerank_not_implemented():
# Create fake model
fake_provider_entity = persistence_model.ModelProvider(
workspace_uuid=TEST_WORKSPACE_UUID,
uuid='provider-uuid',
name='Provider',
requester='test',
@@ -143,17 +145,20 @@ async def test_requester_invoke_rerank_not_implemented():
fake_token_mgr = token.TokenManager(name='test', tokens=[])
fake_requester = inst
fake_provider = requester.RuntimeProvider(
execution_context=TEST_EXECUTION_CONTEXT,
provider_entity=fake_provider_entity,
token_mgr=fake_token_mgr,
requester=fake_requester,
)
fake_model_entity = persistence_model.RerankModel(
workspace_uuid=TEST_WORKSPACE_UUID,
uuid='model-uuid',
name='Model',
provider_uuid='provider-uuid',
extra_args={},
)
fake_model = requester.RuntimeRerankModel(
execution_context=TEST_EXECUTION_CONTEXT,
model_entity=fake_model_entity,
provider=fake_provider,
)
@@ -289,6 +294,7 @@ async def test_runtime_provider_invoke_llm_delegates(runtime_provider, runtime_l
resp_message_chain=None,
current_stage_name=None,
)
object.__setattr__(query, '_execution_context', TEST_EXECUTION_CONTEXT)
messages = [
provider_message.Message(role='user', content=[provider_message.ContentElement(type='text', text='Hello')])
@@ -332,6 +338,7 @@ async def test_runtime_provider_invoke_llm_stream_yields_chunks(runtime_provider
resp_message_chain=None,
current_stage_name=None,
)
object.__setattr__(query, '_execution_context', TEST_EXECUTION_CONTEXT)
messages = [
provider_message.Message(role='user', content=[provider_message.ContentElement(type='text', text='Hello')])
@@ -350,7 +357,11 @@ async def test_runtime_provider_invoke_embedding_returns_vectors(runtime_provide
"""Test RuntimeProvider.invoke_embedding returns embedding vectors."""
provider = runtime_provider
result = await provider.invoke_embedding(runtime_embedding_model, ['text1', 'text2'])
result = await provider.invoke_embedding(
runtime_embedding_model,
['text1', 'text2'],
execution_context=TEST_EXECUTION_CONTEXT,
)
assert len(result) == 2
assert result[0] == [0.1, 0.2, 0.3]
@@ -362,7 +373,12 @@ async def test_runtime_provider_invoke_rerank_returns_scores(runtime_provider, r
# Need to use the correct provider for rerank model
provider = runtime_rerank_model.provider
result = await provider.invoke_rerank(runtime_rerank_model, 'query', ['doc1', 'doc2', 'doc3'])
result = await provider.invoke_rerank(
runtime_rerank_model,
'query',
['doc1', 'doc2', 'doc3'],
execution_context=TEST_EXECUTION_CONTEXT,
)
assert len(result) == 3
assert result[0]['index'] == 0
@@ -532,6 +548,7 @@ async def test_runtime_provider_invoke_llm_propagates_error(mock_app_for_modelmg
await requester_inst.initialize()
provider_entity = persistence_model.ModelProvider(
workspace_uuid=TEST_WORKSPACE_UUID,
uuid='error-provider',
name='Error Provider',
requester='error-requester',
@@ -541,19 +558,25 @@ async def test_runtime_provider_invoke_llm_propagates_error(mock_app_for_modelmg
token_mgr = token.TokenManager(name='error-provider', tokens=['error-key'])
provider = requester.RuntimeProvider(
execution_context=TEST_EXECUTION_CONTEXT,
provider_entity=provider_entity,
token_mgr=token_mgr,
requester=requester_inst,
)
model_entity = persistence_model.LLMModel(
workspace_uuid=TEST_WORKSPACE_UUID,
uuid='error-model',
name='Error Model',
provider_uuid='error-provider',
abilities=[],
extra_args={},
)
model = requester.RuntimeLLMModel(model_entity=model_entity, provider=provider)
model = requester.RuntimeLLMModel(
execution_context=TEST_EXECUTION_CONTEXT,
model_entity=model_entity,
provider=provider,
)
import langbot_plugin.api.entities.builtin.provider.message as provider_message
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
@@ -580,6 +603,7 @@ async def test_runtime_provider_invoke_llm_propagates_error(mock_app_for_modelmg
resp_message_chain=None,
current_stage_name=None,
)
object.__setattr__(query, '_execution_context', TEST_EXECUTION_CONTEXT)
messages = [
provider_message.Message(role='user', content=[provider_message.ContentElement(type='text', text='Hello')])
tests/unit_tests/provider/test_session_manager.py
+273 -4
View File
@@ -10,11 +10,74 @@ from __future__ import annotations
import pytest
import asyncio
from types import SimpleNamespace
from unittest.mock import Mock
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 (
ExecutionContextMismatchError,
ExecutionContextRequiredError,
)
TEST_CONTEXT = ExecutionContext(
instance_uuid='instance-test',
workspace_uuid='workspace-test',
placement_generation=1,
)
TEST_BOT_UUID = 'bot-123'
def bind_query_context(query, *, context=TEST_CONTEXT, bot_uuid=TEST_BOT_UUID):
"""Attach the trusted runtime scope expected by the session manager."""
query.bot_uuid = bot_uuid
query._execution_context = context
return query
def bind_session_context(session, query):
"""Make a mocked legacy Session belong to the Query execution scope."""
session.bot_uuid = query.bot_uuid
session._langbot_session_key = (
TEST_CONTEXT.instance_uuid,
TEST_CONTEXT.workspace_uuid,
TEST_CONTEXT.placement_generation,
query.bot_uuid,
query.launcher_type.value,
query.launcher_id,
)
return session
def scoped_query(
*,
workspace_uuid='workspace-test',
bot_uuid=TEST_BOT_UUID,
placement_generation=1,
pipeline_uuid=None,
):
"""Create a small Query-like object with a complete trusted scope."""
context = ExecutionContext(
instance_uuid='instance-test',
workspace_uuid=workspace_uuid,
placement_generation=placement_generation,
bot_uuid=bot_uuid,
pipeline_uuid=pipeline_uuid,
query_uuid=f'query-{workspace_uuid}-{bot_uuid}',
)
return SimpleNamespace(
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id='same-launcher',
sender_id='same-sender',
bot_uuid=bot_uuid,
_execution_context=context,
)
def get_session_module():
@@ -71,7 +134,7 @@ class TestSessionManagerGetSession:
query.launcher_type = provider_session.LauncherTypes.PERSON
query.launcher_id = '12345'
query.sender_id = '12345'
return query
return bind_query_context(query)
@pytest.mark.asyncio
async def test_creates_new_session_when_not_found(self, mock_app_with_config, sample_query):
@@ -126,11 +189,13 @@ class TestSessionManagerGetSession:
query1.launcher_type = provider_session.LauncherTypes.PERSON
query1.launcher_id = 'user1'
query1.sender_id = 'user1'
bind_query_context(query1)
query2 = Mock(spec=pipeline_query.Query)
query2.launcher_type = provider_session.LauncherTypes.PERSON
query2.launcher_id = 'user2'
query2.sender_id = 'user2'
bind_query_context(query2)
session1 = await manager.get_session(query1)
session2 = await manager.get_session(query2)
@@ -149,11 +214,13 @@ class TestSessionManagerGetSession:
query1.launcher_type = provider_session.LauncherTypes.PERSON
query1.launcher_id = 'same_id'
query1.sender_id = 'same_id'
bind_query_context(query1)
query2 = Mock(spec=pipeline_query.Query)
query2.launcher_type = provider_session.LauncherTypes.GROUP
query2.launcher_id = 'same_id'
query2.sender_id = 'same_id'
bind_query_context(query2)
session1 = await manager.get_session(query1)
session2 = await manager.get_session(query2)
@@ -191,7 +258,7 @@ class TestSessionManagerGetConversation:
query.launcher_type = provider_session.LauncherTypes.PERSON
query.launcher_id = '12345'
query.sender_id = '12345'
return query
return bind_query_context(query)
@pytest.mark.asyncio
async def test_creates_conversation_with_prompt(self, mock_app_with_config, sample_query, sample_session):
@@ -199,6 +266,7 @@ class TestSessionManagerGetConversation:
sessionmgr = get_session_module()
manager = sessionmgr.SessionManager(mock_app_with_config)
bind_session_context(sample_session, sample_query)
prompt_config = [{'role': 'system', 'content': 'You are a helpful assistant.'}]
pipeline_uuid = 'pipeline-123'
@@ -222,6 +290,7 @@ class TestSessionManagerGetConversation:
sessionmgr = get_session_module()
manager = sessionmgr.SessionManager(mock_app_with_config)
bind_session_context(sample_session, sample_query)
prompt_config = [{'role': 'system', 'content': 'You are a helpful assistant.'}]
pipeline_uuid = 'pipeline-123'
@@ -244,14 +313,15 @@ class TestSessionManagerGetConversation:
sessionmgr = get_session_module()
manager = sessionmgr.SessionManager(mock_app_with_config)
bind_session_context(sample_session, sample_query)
prompt_config = [{'role': 'system', 'content': 'You are a helpful assistant.'}]
# First call with pipeline1
conv1 = await manager.get_conversation(sample_query, sample_session, prompt_config, 'pipeline-1', 'bot-1')
conv1 = await manager.get_conversation(sample_query, sample_session, prompt_config, 'pipeline-1', TEST_BOT_UUID)
# Second call with different pipeline should create new conversation
conv2 = await manager.get_conversation(sample_query, sample_session, prompt_config, 'pipeline-2', 'bot-2')
conv2 = await manager.get_conversation(sample_query, sample_session, prompt_config, 'pipeline-2', TEST_BOT_UUID)
assert conv1 is not conv2
assert len(sample_session.conversations) == 2
@@ -263,6 +333,7 @@ class TestSessionManagerGetConversation:
sessionmgr = get_session_module()
manager = sessionmgr.SessionManager(mock_app_with_config)
bind_session_context(sample_session, sample_query)
prompt_config = [{'role': 'system', 'content': 'You are a helpful assistant.'}]
@@ -278,6 +349,7 @@ class TestSessionManagerGetConversation:
sessionmgr = get_session_module()
manager = sessionmgr.SessionManager(mock_app_with_config)
bind_session_context(sample_session, sample_query)
prompt_config = [{'role': 'system', 'content': 'System message'}, {'role': 'user', 'content': 'User message'}]
@@ -287,3 +359,200 @@ class TestSessionManagerGetConversation:
assert conversation.prompt.name == 'default'
assert len(conversation.prompt.messages) == 2
class TestSessionManagerWorkspaceIsolation:
"""Regression coverage for workspace, bot, and placement fencing."""
@staticmethod
def manager():
mock_app = Mock()
mock_app.instance_config.data = {'concurrency': {'session': 5}}
return get_session_module().SessionManager(mock_app)
@pytest.mark.asyncio
async def test_get_session_requires_trusted_query_scope(self):
query = SimpleNamespace(
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id='same-launcher',
sender_id='same-sender',
bot_uuid=TEST_BOT_UUID,
)
with pytest.raises(ExecutionContextRequiredError):
await self.manager().get_session(query)
@pytest.mark.asyncio
async def test_same_launcher_in_two_workspaces_does_not_share_session(self):
manager = self.manager()
first = await manager.get_session(scoped_query(workspace_uuid='workspace-a'))
second = await manager.get_session(scoped_query(workspace_uuid='workspace-b'))
assert first is not second
assert len(manager.session_list) == 2
@pytest.mark.asyncio
async def test_same_launcher_in_two_bots_does_not_share_session(self):
manager = self.manager()
first = await manager.get_session(scoped_query(bot_uuid='bot-a'))
second = await manager.get_session(scoped_query(bot_uuid='bot-b'))
assert first is not second
@pytest.mark.asyncio
async def test_new_placement_generation_does_not_reuse_old_session(self):
manager = self.manager()
first = await manager.get_session(scoped_query(placement_generation=1))
second = await manager.get_session(scoped_query(placement_generation=2))
assert first is not second
@pytest.mark.asyncio
async def test_conversation_rejects_session_from_another_workspace(self):
manager = self.manager()
query_a = scoped_query(workspace_uuid='workspace-a')
query_b = scoped_query(workspace_uuid='workspace-b')
session_a = await manager.get_session(query_a)
with pytest.raises(ExecutionContextMismatchError):
await manager.get_conversation(query_b, session_a, [], 'pipeline-1', TEST_BOT_UUID)
@pytest.mark.asyncio
async def test_conversation_rejects_substituted_bot_argument(self):
manager = self.manager()
query = scoped_query(bot_uuid='bot-a')
session = await manager.get_session(query)
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
tests/unit_tests/provider/test_skill_tools.py
+199 -48
View File
@@ -7,6 +7,37 @@ from unittest.mock import AsyncMock, Mock
import pytest
from langbot.pkg.api.http.context import ExecutionContext
_CONTEXT = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
query_uuid='query-a',
)
def _make_query(*, variables=None, **kwargs):
return SimpleNamespace(
query_id=kwargs.pop('query_id', 'query-a'),
query_uuid=kwargs.pop('query_uuid', 'query-a'),
instance_uuid=kwargs.pop('instance_uuid', _CONTEXT.instance_uuid),
workspace_uuid=kwargs.pop('workspace_uuid', _CONTEXT.workspace_uuid),
placement_generation=kwargs.pop('placement_generation', _CONTEXT.placement_generation),
variables={} if variables is None else variables,
**kwargs,
)
def _make_skill_manager(skills: dict[str, dict], **kwargs):
return SimpleNamespace(
skills=skills,
get_skills=Mock(return_value=skills),
get_skill_by_name=Mock(side_effect=lambda _context, name: skills.get(name)),
**kwargs,
)
def _make_ap(logger=None):
ap = SimpleNamespace()
@@ -51,14 +82,14 @@ class TestSkillManagerCache:
mgr = SkillManager(ap)
# Empty cache → returns False
assert mgr.refresh_skill_from_disk('test-skill') is False
assert mgr.refresh_skill_from_disk(_CONTEXT, 'test-skill') is False
# Cache populated → returns True; method does NOT mutate the cache
cached = _make_skill_data(name='test-skill', instructions='Cached')
mgr.skills['test-skill'] = cached
assert mgr.refresh_skill_from_disk('test-skill') is True
assert mgr.skills['test-skill'] is cached
assert mgr.refresh_skill_from_disk('') is False
mgr._skills_by_scope[mgr._scope_key(_CONTEXT)] = {'test-skill': cached}
assert mgr.refresh_skill_from_disk(_CONTEXT, 'test-skill') is True
assert mgr.get_skills(_CONTEXT)['test-skill'] is cached
assert mgr.refresh_skill_from_disk(_CONTEXT, '') is False
@pytest.mark.asyncio
async def test_reload_skills_drops_box_skills_with_missing_package_root(self):
@@ -85,9 +116,9 @@ class TestSkillManagerCache:
ap.box_service = box_service
mgr = SkillManager(ap)
await mgr.reload_skills()
await mgr.reload_skills(_CONTEXT)
assert list(mgr.skills) == ['alive']
assert list(mgr.get_skills(_CONTEXT)) == ['alive']
# Warning fired with the dropped skill name so operators can see it.
warning_messages = [str(call.args[0]) for call in ap.logger.warning.call_args_list]
assert any('ghost' in msg and 'package_root missing' in msg for msg in warning_messages)
@@ -116,9 +147,9 @@ class TestSkillManagerCache:
ap.box_service = box_service
mgr = SkillManager(ap)
await mgr.reload_skills()
await mgr.reload_skills(_CONTEXT)
assert sorted(mgr.skills) == ['alpha', 'beta']
assert sorted(mgr.get_skills(_CONTEXT)) == ['alpha', 'beta']
# No skill dropped → no "package_root missing" warning.
warning_messages = [str(call.args[0]) for call in ap.logger.warning.call_args_list]
assert not any('package_root missing' in msg for msg in warning_messages)
@@ -141,12 +172,12 @@ class TestSkillActivationHelper:
ap = _make_ap()
mgr = SkillManager(ap)
mgr.skills = {
mgr._skills_by_scope[mgr._scope_key(_CONTEXT)] = {
'primary': _make_skill_data(name='primary', instructions='Primary instructions'),
}
ap.skill_mgr = mgr
query = SimpleNamespace(variables={})
query = _make_query()
assert register_activated_skill(ap, query, 'primary') is True
assert set(query.variables[ACTIVATED_SKILLS_KEY].keys()) == {'primary'}
@@ -159,10 +190,10 @@ class TestSkillActivationHelper:
ap = _make_ap()
mgr = SkillManager(ap)
mgr.skills = {'primary': _make_skill_data(name='primary')}
mgr._skills_by_scope[mgr._scope_key(_CONTEXT)] = {'primary': _make_skill_data(name='primary')}
ap.skill_mgr = mgr
query = SimpleNamespace(variables={})
query = _make_query()
assert register_activated_skill(ap, query, 'missing') is False
assert ACTIVATED_SKILLS_KEY not in query.variables
@@ -171,7 +202,7 @@ class TestSkillActivationHelper:
from langbot.pkg.skill.activation import register_activated_skill
ap = _make_ap() # no skill_mgr attribute
query = SimpleNamespace(variables={})
query = _make_query()
assert register_activated_skill(ap, query, 'primary') is False
@@ -181,13 +212,13 @@ class TestSkillPathHelpers:
from langbot.pkg.provider.tools.loaders.skill import PIPELINE_BOUND_SKILLS_KEY, get_visible_skills
ap = _make_ap()
ap.skill_mgr = SimpleNamespace(
skills={
ap.skill_mgr = _make_skill_manager(
{
'visible': _make_skill_data(name='visible'),
'hidden': _make_skill_data(name='hidden'),
}
)
query = SimpleNamespace(variables={PIPELINE_BOUND_SKILLS_KEY: ['visible']})
query = _make_query(variables={PIPELINE_BOUND_SKILLS_KEY: ['visible']})
result = get_visible_skills(ap, query)
@@ -202,13 +233,13 @@ class TestSkillPathHelpers:
)
ap = _make_ap()
ap.skill_mgr = SimpleNamespace(
skills={
ap.skill_mgr = _make_skill_manager(
{
'visible': _make_skill_data(name='visible'),
'hidden': _make_skill_data(name='hidden'),
}
)
query = SimpleNamespace(variables={PIPELINE_BOUND_SKILLS_KEY: ['visible']})
query = _make_query(variables={PIPELINE_BOUND_SKILLS_KEY: ['visible']})
restored = restore_activated_skills(ap, query, ['visible', 'hidden', 'visible', ''])
@@ -223,8 +254,8 @@ class TestSkillPathHelpers:
)
ap = _make_ap()
ap.skill_mgr = SimpleNamespace(skills={'demo': _make_skill_data(name='demo')})
query = SimpleNamespace(variables={PIPELINE_BOUND_SKILLS_KEY: ['demo']})
ap.skill_mgr = _make_skill_manager({'demo': _make_skill_data(name='demo')})
query = _make_query(variables={PIPELINE_BOUND_SKILLS_KEY: ['demo']})
skill, rewritten = resolve_virtual_skill_path(
ap,
@@ -273,6 +304,22 @@ class TestSkillPathHelpers:
assert 'export VIRTUAL_ENV="$_LB_VENV_DIR"' in command
assert command.rstrip().endswith('python scripts/run.py')
def test_wrap_skill_python_env_keeps_state_outside_read_only_source(self):
from langbot.pkg.provider.tools.loaders.skill import wrap_skill_command_with_python_env
command = wrap_skill_command_with_python_env(
'python scripts/run.py',
mount_path='/workspace/.skills/demo',
state_path='/workspace/.skill-envs/demo',
)
assert '_LB_VENV_DIR="/workspace/.skill-envs/demo/.venv"' in command
assert '_LB_META_DIR="/workspace/.skill-envs/demo/.langbot"' in command
assert '_LB_TMP_DIR="/workspace/.skill-envs/demo/.tmp"' in command
assert '_LB_PIP_CACHE_DIR="/workspace/.skill-envs/demo/.cache/pip"' in command
assert 'root = "/workspace/.skills/demo"' in command
assert 'pip install "/workspace/.skills/demo"' in command
class TestSkillToolLoader:
"""The skill tool surface is now just ``activate`` + ``register_skill``.
@@ -292,13 +339,10 @@ class TestSkillToolLoader:
skill = _make_skill_data(name='demo', package_root='/data/skills/demo', instructions='Step 1')
ap = _make_ap()
ap.skill_mgr = SimpleNamespace(
skills={'demo': skill},
get_skill_by_name=lambda name: skill if name == 'demo' else None,
)
ap.skill_mgr = _make_skill_manager({'demo': skill})
loader = SkillToolLoader(ap)
query = SimpleNamespace(variables={})
query = _make_query()
result = await loader.invoke_tool(ACTIVATE_SKILL_TOOL_NAME, {'skill_name': 'demo'}, query)
@@ -317,10 +361,7 @@ class TestSkillToolLoader:
)
ap = _make_ap()
ap.skill_mgr = SimpleNamespace(
skills={'demo': _make_skill_data(name='demo')},
get_skill_by_name=lambda name: None,
)
ap.skill_mgr = _make_skill_manager({'demo': _make_skill_data(name='demo')})
loader = SkillToolLoader(ap)
@@ -328,7 +369,7 @@ class TestSkillToolLoader:
await loader.invoke_tool(
ACTIVATE_SKILL_TOOL_NAME,
{'skill_name': 'ghost'},
SimpleNamespace(variables={}),
_make_query(),
)
@pytest.mark.asyncio
@@ -362,18 +403,19 @@ class TestSkillToolLoader:
result = await loader.invoke_tool(
REGISTER_SKILL_TOOL_NAME,
{'path': '/workspace/repo'},
SimpleNamespace(),
_make_query(),
)
ap.skill_service.scan_directory_async.assert_awaited_once_with(os.path.realpath(repo_dir))
ap.skill_service.scan_directory_async.assert_awaited_once_with(_CONTEXT, os.path.realpath(repo_dir))
ap.skill_service.create_skill.assert_awaited_once_with(
_CONTEXT,
{
'name': 'cloned-skill',
'display_name': 'Cloned Skill',
'description': 'Imported from clone',
'instructions': 'Do work',
'package_root': os.path.realpath(repo_dir),
}
},
)
assert result['registered'] is True
assert result['skill_name'] == 'cloned-skill'
@@ -397,7 +439,7 @@ class TestSkillToolLoader:
await loader.invoke_tool(
REGISTER_SKILL_TOOL_NAME,
{'path': '/workspace/../../etc'},
SimpleNamespace(),
_make_query(),
)
@pytest.mark.asyncio
@@ -417,7 +459,7 @@ class TestSkillToolLoader:
await loader.invoke_tool(
REGISTER_SKILL_TOOL_NAME,
{'path': '/workspace/foo'},
SimpleNamespace(),
_make_query(),
)
@pytest.mark.asyncio
@@ -428,7 +470,7 @@ class TestSkillToolLoader:
ap.skill_mgr = SimpleNamespace(skills={})
ap.box_service = SimpleNamespace(
available=True,
get_status=AsyncMock(return_value={'backend': {'available': False}}),
get_backend_status=AsyncMock(return_value={'backend': {'available': False}}),
)
loader = SkillToolLoader(ap)
@@ -443,10 +485,10 @@ class TestSkillToolLoader:
from langbot.pkg.provider.tools.loaders.skill_authoring import SkillToolLoader
ap = _make_ap()
ap.skill_mgr = SimpleNamespace(skills={'demo': _make_skill_data(name='demo')})
ap.skill_mgr = _make_skill_manager({'demo': _make_skill_data(name='demo')})
ap.box_service = SimpleNamespace(
available=True,
get_status=AsyncMock(return_value={'backend': {'available': True}}),
get_backend_status=AsyncMock(return_value={'backend': {'available': True}}),
)
loader = SkillToolLoader(ap)
@@ -466,7 +508,7 @@ class TestSkillToolLoader:
ap.skill_mgr = SimpleNamespace(skills={'demo': _make_skill_data(name='demo')})
ap.box_service = SimpleNamespace(
available=False,
get_status=AsyncMock(return_value={'backend': {'available': True}}),
get_backend_status=AsyncMock(return_value={'backend': {'available': True}}),
)
loader = SkillToolLoader(ap)
@@ -490,20 +532,88 @@ class TestNativeToolLoaderSkillPaths:
f.write('demo instructions')
ap = _make_ap()
ap.box_service = SimpleNamespace(available=True, default_workspace=tmpdir)
ap.skill_mgr = SimpleNamespace(skills={'demo': _make_skill_data(name='demo', package_root=tmpdir)})
ap.box_service = SimpleNamespace(
available=True,
default_workspace=tmpdir,
shares_filesystem_with_box=True,
)
ap.skill_mgr = _make_skill_manager({'demo': _make_skill_data(name='demo', package_root=tmpdir)})
loader = NativeToolLoader(ap)
result = await loader.invoke_tool(
'read',
{'path': '/workspace/.skills/demo/SKILL.md'},
SimpleNamespace(query_id='q1', variables={PIPELINE_BOUND_SKILLS_KEY: ['demo']}),
_make_query(query_id='q1', variables={PIPELINE_BOUND_SKILLS_KEY: ['demo']}),
)
assert result['ok'] is True
assert result['content'] == 'demo instructions'
assert result['truncated'] is False
@pytest.mark.asyncio
async def test_external_runtime_read_never_interprets_package_root_on_core_host(self):
from langbot.pkg.provider.tools.loaders.native import NativeToolLoader
from langbot.pkg.provider.tools.loaders.skill import PIPELINE_BOUND_SKILLS_KEY
with tempfile.TemporaryDirectory() as tmpdir:
with open(os.path.join(tmpdir, 'SKILL.md'), 'w', encoding='utf-8') as file_obj:
file_obj.write('core-host-secret')
ap = _make_ap()
ap.box_service = SimpleNamespace(
available=True,
shares_filesystem_with_box=False,
read_skill_file=AsyncMock(return_value={'content': 'runtime-owned-content'}),
)
ap.skill_mgr = _make_skill_manager({'demo': _make_skill_data(name='demo', package_root=tmpdir)})
loader = NativeToolLoader(ap)
query = _make_query(
query_id='q-external-read',
variables={PIPELINE_BOUND_SKILLS_KEY: ['demo']},
)
result = await loader.invoke_tool(
'read',
{'path': '/workspace/.skills/demo/SKILL.md'},
query,
)
assert result['ok'] is True
assert result['content'] == 'runtime-owned-content'
assert 'core-host-secret' not in repr(result)
ap.box_service.read_skill_file.assert_awaited_once_with(_CONTEXT, 'demo', 'SKILL.md')
@pytest.mark.asyncio
async def test_external_runtime_rejects_skill_host_fallback_without_protocol_capability(self):
from langbot.pkg.provider.tools.loaders.native import NativeToolLoader
from langbot.pkg.provider.tools.loaders.skill import PIPELINE_BOUND_SKILLS_KEY
with tempfile.TemporaryDirectory() as tmpdir:
with open(os.path.join(tmpdir, 'secret.txt'), 'w', encoding='utf-8') as file_obj:
file_obj.write('core-host-secret')
ap = _make_ap()
ap.box_service = SimpleNamespace(
available=True,
shares_filesystem_with_box=False,
)
ap.skill_mgr = _make_skill_manager({'demo': _make_skill_data(name='demo', package_root=tmpdir)})
loader = NativeToolLoader(ap)
query = _make_query(
query_id='q-external-no-protocol',
variables={PIPELINE_BOUND_SKILLS_KEY: ['demo']},
)
with pytest.raises(ValueError, match='owned by the Box Runtime'):
await loader.invoke_tool(
'grep',
{
'path': '/workspace/.skills/demo',
'pattern': 'core-host-secret',
},
query,
)
@pytest.mark.asyncio
async def test_exec_in_activated_skill_mount_rewrites_command_and_refreshes(self):
from langbot.pkg.provider.tools.loaders.native import NativeToolLoader
@@ -519,7 +629,7 @@ class TestNativeToolLoaderSkillPaths:
ap.skill_mgr = SimpleNamespace(refresh_skill_from_disk=Mock())
loader = NativeToolLoader(ap)
query = SimpleNamespace(query_id='q1', launcher_type='person', launcher_id='123', variables={})
query = _make_query(query_id='q1', launcher_type='person', launcher_id='123')
register_activated_skill(query, _make_skill_data(name='demo', package_root=tmpdir))
result = await loader.invoke_tool(
@@ -535,7 +645,48 @@ class TestNativeToolLoaderSkillPaths:
tool_parameters = ap.box_service.execute_tool.await_args.args[0]
assert tool_parameters['command'] == 'python /workspace/.skills/demo/scripts/run.py'
assert tool_parameters['workdir'] == '/workspace/.skills/demo'
ap.skill_mgr.refresh_skill_from_disk.assert_called_once_with('demo')
assert ap.box_service.execute_tool.await_args.kwargs['skill_name'] == 'demo'
ap.skill_mgr.refresh_skill_from_disk.assert_called_once_with(_CONTEXT, 'demo')
@pytest.mark.asyncio
async def test_external_runtime_python_skill_uses_trusted_metadata_and_writable_env(self):
from langbot.pkg.provider.tools.loaders.native import NativeToolLoader
from langbot.pkg.provider.tools.loaders.skill import register_activated_skill
ap = _make_ap()
ap.box_service = SimpleNamespace(
available=True,
shares_filesystem_with_box=False,
execute_tool=AsyncMock(return_value={'ok': True}),
)
ap.skill_mgr = SimpleNamespace(refresh_skill_from_disk=Mock())
loader = NativeToolLoader(ap)
query = _make_query(query_id='q-external', launcher_type='person', launcher_id='123')
register_activated_skill(
query,
_make_skill_data(
name='demo',
package_root='/box-runtime/skills/tenants/workspace/demo',
python_project=True,
),
)
result = await loader.invoke_tool(
'exec',
{
'command': 'python /workspace/.skills/demo/scripts/run.py',
'workdir': '/workspace/.skills/demo',
},
query,
)
assert result['ok'] is True
tool_parameters = ap.box_service.execute_tool.await_args.args[0]
wrapped = tool_parameters['command']
assert '_LB_VENV_DIR="/workspace/.skill-envs/demo/.venv"' in wrapped
assert 'root = "/workspace/.skills/demo"' in wrapped
assert '/box-runtime/skills/tenants/workspace/demo' not in wrapped
assert ap.box_service.execute_tool.await_args.kwargs['skill_name'] == 'demo'
@pytest.mark.asyncio
async def test_write_requires_skill_activation(self):
@@ -545,10 +696,10 @@ class TestNativeToolLoaderSkillPaths:
with tempfile.TemporaryDirectory() as tmpdir:
ap = _make_ap()
ap.box_service = SimpleNamespace(available=True, default_workspace=tmpdir)
ap.skill_mgr = SimpleNamespace(skills={'demo': _make_skill_data(name='demo', package_root=tmpdir)})
ap.skill_mgr = _make_skill_manager({'demo': _make_skill_data(name='demo', package_root=tmpdir)})
loader = NativeToolLoader(ap)
query = SimpleNamespace(query_id='q1', variables={PIPELINE_BOUND_SKILLS_KEY: ['demo']})
query = _make_query(query_id='q1', variables={PIPELINE_BOUND_SKILLS_KEY: ['demo']})
with pytest.raises(ValueError, match='Skill "demo" is not available at this path'):
await loader.invoke_tool(
tests/unit_tests/provider/test_tool_manager.py
+13
View File
@@ -14,6 +14,15 @@ from importlib import import_module
import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
from langbot.pkg.api.http.context import ExecutionContext
_CONTEXT = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
def get_toolmgr_module():
"""Lazy import to avoid circular import issues."""
@@ -188,6 +197,10 @@ class TestToolManagerExecuteFuncCall:
def sample_query(self):
"""Create sample query for testing."""
query = Mock(spec=pipeline_query.Query)
query.bot_uuid = None
query.pipeline_uuid = None
query.query_uuid = None
query._execution_context = _CONTEXT
return query
@pytest.mark.asyncio
tests/unit_tests/provider/test_tool_manager_native.py
+296 -14
View File
@@ -1,8 +1,10 @@
from __future__ import annotations
import base64
import contextlib
import os
import tempfile
import threading
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
@@ -11,7 +13,16 @@ import pytest
import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
from langbot.pkg.provider.tools.loaders.native import NativeToolLoader
from langbot.pkg.provider.tools.loaders import native as native_loader
from langbot.pkg.provider.tools.toolmgr import ToolManager
from langbot.pkg.api.http.context import ExecutionContext
_CONTEXT = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
class StubLoader:
@@ -44,7 +55,8 @@ class StubLoader:
for tool in self._tools
]
async def has_tool(self, name: str) -> bool:
async def has_tool(self, *args) -> bool:
name = args[-1]
return any(tool.name == name for tool in self._tools)
async def invoke_tool(self, name: str, parameters: dict, query):
@@ -72,7 +84,7 @@ async def test_tool_manager_omits_skill_authoring_tools_by_default():
manager.plugin_tool_loader = StubLoader([make_tool('plugin_tool')])
manager.mcp_tool_loader = StubLoader([make_tool('mcp_tool')])
tools = await manager.get_all_tools()
tools = await manager.get_all_tools(_CONTEXT)
assert [tool.name for tool in tools] == ['exec', 'plugin_tool', 'mcp_tool']
@@ -85,7 +97,7 @@ async def test_tool_manager_includes_skill_authoring_tools_when_requested():
manager.plugin_tool_loader = StubLoader([make_tool('plugin_tool')])
manager.mcp_tool_loader = StubLoader([make_tool('mcp_tool')])
tools = await manager.get_all_tools(include_skill_authoring=True)
tools = await manager.get_all_tools(_CONTEXT, include_skill_authoring=True)
assert [tool.name for tool in tools] == ['exec', 'activate', 'plugin_tool', 'mcp_tool']
@@ -102,7 +114,7 @@ async def test_tool_manager_catalog_labels_tool_sources():
)
manager.mcp_tool_loader = StubLoader([make_tool('mcp_tool')])
catalog = await manager.get_tool_catalog(include_skill_authoring=True)
catalog = await manager.get_tool_catalog(_CONTEXT, include_skill_authoring=True)
assert [(item['name'], item['source'], item['source_name']) for item in catalog] == [
('exec', 'builtin', 'LangBot'),
@@ -121,11 +133,55 @@ async def test_tool_manager_routes_native_tool_calls():
manager.plugin_tool_loader = StubLoader([make_tool('plugin_tool')])
manager.mcp_tool_loader = StubLoader([make_tool('mcp_tool')])
result = await manager.execute_func_call('exec', {'command': 'pwd'}, query=Mock())
query = SimpleNamespace(
_execution_context=_CONTEXT,
bot_uuid=None,
pipeline_uuid=None,
query_uuid=None,
)
result = await manager.execute_func_call('exec', {'command': 'pwd'}, query=query)
assert result == {'backend': 'fake'}
@pytest.mark.asyncio
async def test_tool_manager_hides_sandbox_and_skill_tools_without_workspace_entitlement():
box_service = SimpleNamespace(is_workspace_sandbox_available=AsyncMock(return_value=False))
manager = ToolManager(SimpleNamespace(box_service=box_service))
manager.native_tool_loader = StubLoader([make_tool('exec')])
manager.skill_tool_loader = StubLoader([make_tool('activate')])
manager.plugin_tool_loader = StubLoader([make_tool('plugin_tool')])
manager.mcp_tool_loader = StubLoader([make_tool('mcp_tool')])
tools = await manager.get_all_tools(_CONTEXT, include_skill_authoring=True)
catalog = await manager.get_tool_catalog(_CONTEXT, include_skill_authoring=True)
assert [tool.name for tool in tools] == ['plugin_tool', 'mcp_tool']
assert [item['name'] for item in catalog] == ['plugin_tool', 'mcp_tool']
assert box_service.is_workspace_sandbox_available.await_count == 2
@pytest.mark.asyncio
async def test_tool_manager_rechecks_workspace_entitlement_before_native_invocation():
box_service = SimpleNamespace(is_workspace_sandbox_available=AsyncMock(return_value=False))
manager = ToolManager(SimpleNamespace(box_service=box_service))
manager.native_tool_loader = StubLoader([make_tool('exec')], invoke_result={'unexpected': True})
manager.skill_tool_loader = StubLoader([])
manager.plugin_tool_loader = StubLoader([])
manager.mcp_tool_loader = StubLoader([])
query = SimpleNamespace(
_execution_context=_CONTEXT,
bot_uuid=None,
pipeline_uuid=None,
query_uuid=None,
)
with pytest.raises(Exception, match='exec'):
await manager.execute_func_call('exec', {'command': 'pwd'}, query=query)
box_service.is_workspace_sandbox_available.assert_awaited_once_with(_CONTEXT)
@pytest.mark.asyncio
async def test_native_tool_loader_hides_tools_when_box_unavailable():
loader = NativeToolLoader(SimpleNamespace(box_service=SimpleNamespace(available=False)))
@@ -139,7 +195,7 @@ async def test_native_tool_loader_hides_tools_when_box_unavailable():
async def test_native_tool_loader_exposes_all_tools_when_box_available():
box_service = SimpleNamespace(
available=True,
get_status=AsyncMock(return_value={'backend': {'available': True}}),
get_backend_status=AsyncMock(return_value={'backend': {'available': True}}),
)
loader = NativeToolLoader(SimpleNamespace(box_service=box_service, logger=Mock()))
await loader.initialize()
@@ -155,7 +211,7 @@ async def test_native_tool_loader_exposes_all_tools_when_box_available():
async def test_native_tool_loader_refreshes_after_box_recovers():
box_service = SimpleNamespace(
available=False,
get_status=AsyncMock(return_value={'backend': {'available': True}}),
get_backend_status=AsyncMock(return_value={'backend': {'available': True}}),
)
loader = NativeToolLoader(SimpleNamespace(box_service=box_service, logger=Mock()))
await loader.initialize()
@@ -166,20 +222,51 @@ async def test_native_tool_loader_refreshes_after_box_recovers():
assert [tool.name for tool in await loader.get_tools()] == ['exec', 'read', 'write', 'edit', 'glob', 'grep']
@pytest.mark.asyncio
async def test_native_tool_loader_rechecks_admission_at_the_final_invoke_boundary():
box_service = SimpleNamespace(
available=True,
require_workspace_sandbox=AsyncMock(side_effect=RuntimeError('entitlement expired')),
)
loader = NativeToolLoader(SimpleNamespace(box_service=box_service, logger=Mock()))
query = SimpleNamespace(
_execution_context=_CONTEXT,
bot_uuid=None,
pipeline_uuid=None,
query_uuid=None,
)
with pytest.raises(RuntimeError, match='entitlement expired'):
await loader.invoke_tool('read', {'path': '/workspace/private.txt'}, query)
box_service.require_workspace_sandbox.assert_awaited_once_with(_CONTEXT)
# ── read/write/edit file tool tests ─────────────────────────────
def _make_loader_with_workspace(tmpdir: str) -> tuple[NativeToolLoader, Mock]:
logger = Mock()
box_service = SimpleNamespace(available=True, default_workspace=tmpdir)
box_service = SimpleNamespace(
available=True,
default_workspace=tmpdir,
_tenant_workspace=Mock(return_value=tmpdir),
)
ap = SimpleNamespace(box_service=box_service, logger=logger)
return NativeToolLoader(ap), logger
def _make_query() -> Mock:
q = Mock()
q.query_id = 'test-query-1'
return q
def _make_query() -> SimpleNamespace:
return SimpleNamespace(
query_id='test-query-1',
query_uuid='test-query-1',
instance_uuid=_CONTEXT.instance_uuid,
workspace_uuid=_CONTEXT.workspace_uuid,
placement_generation=_CONTEXT.placement_generation,
bot_uuid=None,
pipeline_uuid=None,
variables={},
)
@pytest.mark.asyncio
@@ -373,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:
@@ -382,6 +487,137 @@ async def test_path_escape_blocked():
await loader.invoke_tool('read', {'path': '/workspace/../../etc/passwd'}, _make_query())
@pytest.mark.parametrize(
('tool_name', 'parameters'),
[
('read', {'path': '/workspace/shared/tenant-b-only.txt'}),
(
'write',
{'path': '/workspace/shared/tenant-b-only.txt', 'content': 'overwritten by tenant a'},
),
(
'edit',
{
'path': '/workspace/shared/tenant-b-only.txt',
'old_string': 'tenant-b-secret',
'new_string': 'overwritten by tenant a',
},
),
('glob', {'path': '/workspace/shared', 'pattern': '*'}),
('grep', {'path': '/workspace/shared', 'pattern': 'tenant-b-secret'}),
],
)
@pytest.mark.asyncio
async def test_host_workspace_operations_do_not_follow_a_swapped_ancestor(
monkeypatch,
tool_name: str,
parameters: dict,
):
with tempfile.TemporaryDirectory() as tmpdir:
tenant_a = os.path.join(tmpdir, 'tenant-a')
tenant_b = os.path.join(tmpdir, 'tenant-b')
os.makedirs(os.path.join(tenant_a, 'shared'))
os.makedirs(os.path.join(tenant_b, 'shared'))
tenant_b_file = os.path.join(tenant_b, 'shared', 'tenant-b-only.txt')
with open(os.path.join(tenant_a, 'shared', 'tenant-a-only.txt'), 'w', encoding='utf-8') as file_obj:
file_obj.write('tenant-a-content')
with open(tenant_b_file, 'w', encoding='utf-8') as file_obj:
file_obj.write('tenant-b-secret')
box_service = SimpleNamespace(
available=True,
default_workspace=tmpdir,
_tenant_workspace=Mock(return_value=tenant_a),
)
loader = NativeToolLoader(SimpleNamespace(box_service=box_service, logger=Mock()))
original_open_host_root = native_loader._open_host_root
@contextlib.contextmanager
def open_host_root_after_swap(location, *, create):
with original_open_host_root(location, create=create) as root_fd:
original_ancestor = os.path.join(tenant_a, 'shared-original')
os.rename(os.path.join(tenant_a, 'shared'), original_ancestor)
os.symlink(os.path.join(tenant_b, 'shared'), os.path.join(tenant_a, 'shared'))
yield root_fd
monkeypatch.setattr(native_loader, '_open_host_root', open_host_root_after_swap)
try:
result = await loader.invoke_tool(tool_name, parameters, _make_query())
except ValueError as exc:
result = {'ok': False, 'error': str(exc)}
assert result.get('ok') is False
assert 'tenant-b-secret' not in repr(result)
assert 'tenant-b-only.txt' not in repr(result)
with open(tenant_b_file, encoding='utf-8') as file_obj:
assert file_obj.read() == 'tenant-b-secret'
@pytest.mark.asyncio
async def test_host_file_api_falls_back_to_tenant_box_when_openat_is_unavailable(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": true, "content": "box-owned", "truncated": false}',
'stderr': '',
}
),
)
loader = NativeToolLoader(SimpleNamespace(box_service=box_service, logger=Mock()))
monkeypatch.setattr(native_loader, '_SECURE_HOST_FILE_OPS_AVAILABLE', False)
result = await loader.invoke_tool(
'read',
{'path': '/workspace/file.txt'},
_make_query(),
)
assert result['ok'] is True
assert result['content'] == 'box-owned'
command = box_service.execute_tool.await_args.args[0]['command']
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
@@ -391,13 +627,13 @@ async def test_box_availability_helper_handles_unavailable_and_errors():
unavailable_backend = SimpleNamespace(
available=True,
get_status=AsyncMock(return_value={'backend': {'available': False}}),
get_backend_status=AsyncMock(return_value={'backend': {'available': False}}),
)
assert await is_box_backend_available(SimpleNamespace(box_service=unavailable_backend)) is False
failing_backend = SimpleNamespace(
available=True,
get_status=AsyncMock(side_effect=RuntimeError('box unavailable')),
get_backend_status=AsyncMock(side_effect=RuntimeError('box unavailable')),
)
assert await is_box_backend_available(SimpleNamespace(box_service=failing_backend)) is False
@@ -495,6 +731,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:
@@ -512,3 +776,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'}