xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-07-22 04:16:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(tenancy): implement workspace isolation

Browse Source
This commit is contained in:
Junyan Qin
2026-07-19 09:58:59 +08:00
parent 9eb292992d
commit c6f826fe2d
271 changed files with 31162 additions and 6106 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
+24 -1
View File
@@ -252,42 +252,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'
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
+68 -8
View File
@@ -154,7 +154,16 @@ def mcp_module():
def _make_ap():
ap = Mock()
ap.logger = Mock()
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 +175,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,
),
)
@@ -417,7 +431,7 @@ class TestBuildBoxSessionPayload:
payload = s._build_box_session_payload('session-123')
assert payload['image'] == 'node:20'
assert payload['cpus'] == 2.0
assert payload["memory_mb"] == 1024
assert payload['memory_mb'] == 1024
assert payload['pids_limit'] == 256
def test_none_fields_excluded(self, mcp_module):
@@ -591,6 +605,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,
@@ -774,12 +808,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
@@ -790,7 +828,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()
@@ -814,7 +862,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
@@ -824,11 +873,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
@@ -867,7 +927,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,
@@ -933,7 +993,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
+9
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 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
@@ -133,6 +141,7 @@ def _session(url: str, *, timeout: float = 2) -> RuntimeMCPSession:
{'uuid': 'srv-1', 'mode': 'remote', 'url': url, 'timeout': timeout},
True,
app,
TEST_EXECUTION_CONTEXT,
)
tests/unit_tests/provider/test_mcp_resources.py
+149 -13
View File
@@ -8,6 +8,7 @@ import httpx
import pytest
from mcp import types as mcp_types
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,
@@ -18,10 +19,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(
@@ -30,8 +51,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 [
@@ -51,8 +79,20 @@ 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.sessions[loader._session_key(session.execution_context, session.server_name)] = session
def _http_status_error(status_code: int) -> httpx.HTTPStatusError:
@@ -68,6 +108,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)])
@@ -87,6 +128,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)
@@ -244,10 +286,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,
@@ -260,9 +310,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,
}
@@ -301,9 +351,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'},
@@ -321,3 +372,88 @@ async def test_build_resource_context_for_query_uses_only_bound_attached_text_re
assert query.variables[MCP_RESOURCE_CONTEXT_QUERY_KEY]['resource_count'] == 1
docs.session.read_resource.assert_awaited_once()
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', {})
@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_mcp_idle_lifecycle_stops_without_retry_after_generation_bump():
session = _connected_session()
session.server_config.update({'mode': 'remote', 'url': 'https://example.com/mcp'})
session._FENCE_POLL_INTERVAL = 0
session._init_remote_server = AsyncMock()
session.refresh = AsyncMock()
session._assert_execution_active = AsyncMock(
side_effect=[
None,
None,
None,
WorkspaceGenerationMismatchError('generation changed while idle'),
]
)
await session._lifecycle_loop_with_retry()
session._init_remote_server.assert_awaited_once_with()
assert session.status == MCPSessionStatus.ERROR
assert session.error_message == 'Workspace execution binding is stale'
assert session._shutdown_event.is_set()
tests/unit_tests/provider/test_model_manager.py
+179 -45
View File
@@ -8,14 +8,22 @@ and error handling without calling real LLM APIs.
from __future__ import annotations
import pytest
from unittest.mock import Mock
from unittest.mock import AsyncMock, Mock
from langbot.pkg.provider.modelmgr.modelmgr import ModelManager
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
from tests.unit_tests.provider.conftest import (
TEST_EXECUTION_CONTEXT,
TEST_WORKSPACE_UUID,
_make_mock_result,
_make_row_mock,
)
# ============================================================================
@@ -91,13 +99,15 @@ 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
@@ -118,7 +128,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'
@@ -137,7 +147,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'
@@ -154,7 +164,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
@@ -181,7 +191,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'
@@ -194,7 +204,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)
@@ -215,7 +225,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'
@@ -227,7 +240,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
@@ -246,7 +262,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'
@@ -258,7 +274,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')
# ============================================================================
@@ -282,12 +298,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
@@ -306,12 +322,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
@@ -330,11 +346,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
@@ -353,11 +369,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
@@ -376,11 +392,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()
)
# ============================================================================
@@ -498,7 +520,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'
@@ -528,7 +550,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'
@@ -553,7 +578,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'
@@ -589,12 +617,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',
@@ -605,9 +637,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
@@ -624,7 +659,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'
@@ -643,7 +678,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
@@ -659,7 +698,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
@@ -673,7 +716,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
@@ -692,6 +739,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,
@@ -699,7 +747,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
@@ -723,6 +775,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',
@@ -736,7 +789,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()
@@ -750,6 +803,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',
@@ -759,6 +813,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',
@@ -773,7 +828,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()
@@ -790,6 +845,85 @@ 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_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
+143 -4
View File
@@ -10,12 +10,73 @@ 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
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():
"""Lazy import to avoid circular import issues."""
@@ -71,7 +132,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 +187,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 +212,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 +256,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 +264,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 +288,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 +311,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 +331,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 +347,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 +357,72 @@ 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')
tests/unit_tests/provider/test_skill_tools.py
+72 -46
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,
@@ -292,13 +323,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 +345,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 +353,7 @@ class TestSkillToolLoader:
await loader.invoke_tool(
ACTIVATE_SKILL_TOOL_NAME,
{'skill_name': 'ghost'},
SimpleNamespace(variables={}),
_make_query(),
)
@pytest.mark.asyncio
@@ -362,18 +387,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 +423,7 @@ class TestSkillToolLoader:
await loader.invoke_tool(
REGISTER_SKILL_TOOL_NAME,
{'path': '/workspace/../../etc'},
SimpleNamespace(),
_make_query(),
)
@pytest.mark.asyncio
@@ -417,7 +443,7 @@ class TestSkillToolLoader:
await loader.invoke_tool(
REGISTER_SKILL_TOOL_NAME,
{'path': '/workspace/foo'},
SimpleNamespace(),
_make_query(),
)
@pytest.mark.asyncio
@@ -428,7 +454,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 +469,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)
@@ -472,13 +498,13 @@ class TestNativeToolLoaderSkillPaths:
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)
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
@@ -500,7 +526,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(
@@ -516,7 +542,7 @@ 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')
ap.skill_mgr.refresh_skill_from_disk.assert_called_once_with(_CONTEXT, 'demo')
@pytest.mark.asyncio
async def test_write_requires_skill_activation(self):
@@ -526,10 +552,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
+30 -11
View File
@@ -12,6 +12,14 @@ 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.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:
@@ -72,7 +80,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 +93,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 +110,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'),
@@ -139,7 +147,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()
@@ -156,15 +164,26 @@ async def test_native_tool_loader_exposes_all_tools_when_box_available():
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
@@ -376,13 +395,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