mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-07-15 08:56:07 +00:00
Merge remote-tracking branch 'origin/master' into feat/card_dify_human_input
This commit is contained in:
fdc310
@@ -1 +1 @@
|
||||
|
||||
"""Provider requester tests"""
|
||||
|
||||
@@ -88,7 +88,10 @@ class AnotherFakeRequester(requester.ProviderAPIRequester):
|
||||
|
||||
async def invoke_llm(self, query, model, messages, funcs=None, extra_args={}, remove_think=False):
|
||||
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
||||
return provider_message.Message(role='assistant', content=[provider_message.ContentElement(type='text', text='Another response')])
|
||||
|
||||
return provider_message.Message(
|
||||
role='assistant', content=[provider_message.ContentElement(type='text', text='Another response')]
|
||||
)
|
||||
|
||||
async def invoke_rerank(self, model, query: str, documents: list, extra_args={}):
|
||||
"""Return fake rerank results."""
|
||||
@@ -135,8 +138,10 @@ def mock_app_for_modelmgr():
|
||||
|
||||
# Fake persistence manager - returns empty results by default
|
||||
app.persistence_mgr = SimpleNamespace()
|
||||
|
||||
async def default_execute(query):
|
||||
return _make_mock_result([])
|
||||
|
||||
app.persistence_mgr.execute_async = AsyncMock(side_effect=default_execute)
|
||||
|
||||
# Fake discover engine
|
||||
@@ -165,9 +170,7 @@ def fake_requester_registry(mock_app_for_modelmgr):
|
||||
fake_component = _create_fake_component('fake-requester', FakeProviderAPIRequester)
|
||||
another_component = _create_fake_component('another-fake-requester', AnotherFakeRequester)
|
||||
|
||||
app.discover.get_components_by_kind = Mock(
|
||||
return_value=[fake_component, another_component]
|
||||
)
|
||||
app.discover.get_components_by_kind = Mock(return_value=[fake_component, another_component])
|
||||
|
||||
model_mgr = ModelManager(app)
|
||||
return model_mgr
|
||||
|
||||
@@ -1,32 +0,0 @@
|
||||
"""Tests for AnthropicMessages requester.
|
||||
|
||||
Tests config and pure utility methods.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
|
||||
class TestAnthropicMessagesConfig:
|
||||
"""Tests for default config."""
|
||||
|
||||
def test_default_config_values(self):
|
||||
"""Check default_config."""
|
||||
from langbot.pkg.provider.modelmgr.requesters.anthropicmsgs import AnthropicMessages
|
||||
|
||||
assert AnthropicMessages.default_config['base_url'] == 'https://api.anthropic.com'
|
||||
assert AnthropicMessages.default_config['timeout'] == 120
|
||||
|
||||
def test_config_override(self):
|
||||
"""Config can override defaults."""
|
||||
from langbot.pkg.provider.modelmgr.requesters.anthropicmsgs import AnthropicMessages
|
||||
|
||||
mock_app = MagicMock()
|
||||
req = AnthropicMessages(mock_app, {
|
||||
'base_url': 'https://custom.anthropic.com',
|
||||
'timeout': 60,
|
||||
})
|
||||
|
||||
assert req.requester_cfg['base_url'] == 'https://custom.anthropic.com'
|
||||
assert req.requester_cfg['timeout'] == 60
|
||||
@@ -1,247 +0,0 @@
|
||||
"""Tests for requester error handling - direct import version.
|
||||
|
||||
Tests error handling branches by importing real packages and mocking
|
||||
only the necessary dependencies.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
import pytest
|
||||
import openai # Import real openai package
|
||||
|
||||
from langbot.pkg.provider.modelmgr.errors import RequesterError
|
||||
|
||||
|
||||
class TestInvokeLLMErrorHandling:
|
||||
"""Tests for invoke_llm error handling branches."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_app(self):
|
||||
"""Create mock Application."""
|
||||
app = MagicMock()
|
||||
app.tool_mgr = MagicMock()
|
||||
app.tool_mgr.generate_tools_for_openai = AsyncMock(return_value=[])
|
||||
return app
|
||||
|
||||
@pytest.fixture
|
||||
def mock_model(self):
|
||||
"""Create mock RuntimeLLMModel."""
|
||||
model = MagicMock()
|
||||
model.model_entity = MagicMock()
|
||||
model.model_entity.name = 'gpt-4'
|
||||
model.provider = MagicMock()
|
||||
model.provider.token_mgr = MagicMock()
|
||||
model.provider.token_mgr.get_token = MagicMock(return_value='test-key')
|
||||
return model
|
||||
|
||||
@pytest.fixture
|
||||
def mock_message(self):
|
||||
"""Create mock provider message."""
|
||||
msg = MagicMock()
|
||||
msg.dict = MagicMock(return_value={'role': 'user', 'content': 'test'})
|
||||
return msg
|
||||
|
||||
@pytest.fixture
|
||||
def requester_with_mocked_client(self, mock_app):
|
||||
"""Create requester with mocked OpenAI client."""
|
||||
from langbot.pkg.provider.modelmgr.requesters.chatcmpl import OpenAIChatCompletions
|
||||
|
||||
req = OpenAIChatCompletions(mock_app, {
|
||||
'base_url': 'https://api.openai.com/v1',
|
||||
'timeout': 120,
|
||||
})
|
||||
|
||||
# Replace client with mock
|
||||
req.client = MagicMock()
|
||||
req.client.chat = MagicMock()
|
||||
req.client.chat.completions = MagicMock()
|
||||
req.client.chat.completions.create = AsyncMock()
|
||||
|
||||
return req
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_timeout_error(self, requester_with_mocked_client, mock_model, mock_message):
|
||||
"""TimeoutError is wrapped as RequesterError."""
|
||||
requester_with_mocked_client.client.chat.completions.create = AsyncMock(
|
||||
side_effect=asyncio.TimeoutError()
|
||||
)
|
||||
|
||||
with pytest.raises(RequesterError) as exc:
|
||||
await requester_with_mocked_client.invoke_llm(
|
||||
query=None,
|
||||
model=mock_model,
|
||||
messages=[mock_message],
|
||||
)
|
||||
|
||||
assert '超时' in str(exc.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bad_request_context_length(self, requester_with_mocked_client, mock_model, mock_message):
|
||||
"""BadRequestError with context_length_exceeded has special message."""
|
||||
error = openai.BadRequestError(
|
||||
message='context_length_exceeded: max 4096',
|
||||
response=MagicMock(status_code=400),
|
||||
body={}
|
||||
)
|
||||
requester_with_mocked_client.client.chat.completions.create = AsyncMock(
|
||||
side_effect=error
|
||||
)
|
||||
|
||||
with pytest.raises(RequesterError) as exc:
|
||||
await requester_with_mocked_client.invoke_llm(
|
||||
query=None,
|
||||
model=mock_model,
|
||||
messages=[mock_message],
|
||||
)
|
||||
|
||||
assert '上文过长' in str(exc.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authentication_error(self, requester_with_mocked_client, mock_model, mock_message):
|
||||
"""AuthenticationError shows invalid api-key message."""
|
||||
error = openai.AuthenticationError(
|
||||
message='Invalid API key',
|
||||
response=MagicMock(status_code=401),
|
||||
body={}
|
||||
)
|
||||
requester_with_mocked_client.client.chat.completions.create = AsyncMock(
|
||||
side_effect=error
|
||||
)
|
||||
|
||||
with pytest.raises(RequesterError) as exc:
|
||||
await requester_with_mocked_client.invoke_llm(
|
||||
query=None,
|
||||
model=mock_model,
|
||||
messages=[mock_message],
|
||||
)
|
||||
|
||||
assert 'api-key' in str(exc.value).lower() or '无效' in str(exc.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rate_limit_error(self, requester_with_mocked_client, mock_model, mock_message):
|
||||
"""RateLimitError shows rate limit message."""
|
||||
error = openai.RateLimitError(
|
||||
message='Rate limit exceeded',
|
||||
response=MagicMock(status_code=429),
|
||||
body={}
|
||||
)
|
||||
requester_with_mocked_client.client.chat.completions.create = AsyncMock(
|
||||
side_effect=error
|
||||
)
|
||||
|
||||
with pytest.raises(RequesterError) as exc:
|
||||
await requester_with_mocked_client.invoke_llm(
|
||||
query=None,
|
||||
model=mock_model,
|
||||
messages=[mock_message],
|
||||
)
|
||||
|
||||
assert '频繁' in str(exc.value) or '余额' in str(exc.value)
|
||||
|
||||
|
||||
class TestInvokeEmbeddingErrorHandling:
|
||||
"""Tests for invoke_embedding error handling."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_app(self):
|
||||
return MagicMock()
|
||||
|
||||
@pytest.fixture
|
||||
def mock_embedding_model(self):
|
||||
model = MagicMock()
|
||||
model.model_entity = MagicMock()
|
||||
model.model_entity.name = 'text-embedding-ada-002'
|
||||
model.model_entity.extra_args = {}
|
||||
model.provider = MagicMock()
|
||||
model.provider.token_mgr = MagicMock()
|
||||
model.provider.token_mgr.get_token = MagicMock(return_value='test-key')
|
||||
return model
|
||||
|
||||
@pytest.fixture
|
||||
def requester_with_mocked_client(self, mock_app):
|
||||
from langbot.pkg.provider.modelmgr.requesters.chatcmpl import OpenAIChatCompletions
|
||||
|
||||
req = OpenAIChatCompletions(mock_app, {})
|
||||
req.client = MagicMock()
|
||||
req.client.embeddings = MagicMock()
|
||||
req.client.embeddings.create = AsyncMock()
|
||||
|
||||
return req
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_embedding_timeout_error(self, requester_with_mocked_client, mock_embedding_model):
|
||||
"""TimeoutError in embedding request."""
|
||||
requester_with_mocked_client.client.embeddings.create = AsyncMock(
|
||||
side_effect=asyncio.TimeoutError()
|
||||
)
|
||||
|
||||
with pytest.raises(RequesterError) as exc:
|
||||
await requester_with_mocked_client.invoke_embedding(
|
||||
model=mock_embedding_model,
|
||||
input_text=['test'],
|
||||
)
|
||||
|
||||
assert '超时' in str(exc.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_embedding_bad_request_error(self, requester_with_mocked_client, mock_embedding_model):
|
||||
"""BadRequestError in embedding request."""
|
||||
error = openai.BadRequestError(
|
||||
message='Invalid model',
|
||||
response=MagicMock(status_code=400),
|
||||
body={}
|
||||
)
|
||||
requester_with_mocked_client.client.embeddings.create = AsyncMock(
|
||||
side_effect=error
|
||||
)
|
||||
|
||||
with pytest.raises(RequesterError) as exc:
|
||||
await requester_with_mocked_client.invoke_embedding(
|
||||
model=mock_embedding_model,
|
||||
input_text=['test'],
|
||||
)
|
||||
|
||||
assert '参数' in str(exc.value)
|
||||
|
||||
|
||||
class TestRequesterErrorClass:
|
||||
"""Tests for RequesterError."""
|
||||
|
||||
def test_error_message_prefix(self):
|
||||
"""RequesterError has '模型请求失败' prefix."""
|
||||
from langbot.pkg.provider.modelmgr.errors import RequesterError
|
||||
|
||||
error = RequesterError('test error')
|
||||
assert '模型请求失败' in str(error)
|
||||
|
||||
def test_error_is_exception(self):
|
||||
"""RequesterError inherits Exception."""
|
||||
from langbot.pkg.provider.modelmgr.errors import RequesterError
|
||||
|
||||
error = RequesterError('test')
|
||||
assert isinstance(error, Exception)
|
||||
|
||||
|
||||
class TestDefaultConfig:
|
||||
"""Tests for requester default config."""
|
||||
|
||||
def test_default_config(self):
|
||||
"""Check default_config values."""
|
||||
from langbot.pkg.provider.modelmgr.requesters.chatcmpl import OpenAIChatCompletions
|
||||
|
||||
assert OpenAIChatCompletions.default_config['base_url'] == 'https://api.openai.com/v1'
|
||||
assert OpenAIChatCompletions.default_config['timeout'] == 120
|
||||
|
||||
def test_config_override(self):
|
||||
"""Config overrides defaults."""
|
||||
from langbot.pkg.provider.modelmgr.requesters.chatcmpl import OpenAIChatCompletions
|
||||
|
||||
req = OpenAIChatCompletions(MagicMock(), {
|
||||
'base_url': 'https://custom.com/v1',
|
||||
'timeout': 60,
|
||||
})
|
||||
|
||||
assert req.requester_cfg['base_url'] == 'https://custom.com/v1'
|
||||
assert req.requester_cfg['timeout'] == 60
|
||||
@@ -1,340 +0,0 @@
|
||||
"""Tests for requester pure utility functions.
|
||||
|
||||
Tests the helper methods in OpenAIChatCompletions that don't require network calls.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from tests.utils.import_isolation import isolated_sys_modules
|
||||
|
||||
|
||||
class TestMaskApiKey:
|
||||
"""Tests for _mask_api_key method."""
|
||||
|
||||
def _create_requester_with_mocks(self):
|
||||
"""Create requester instance with mocked dependencies."""
|
||||
mocks = {
|
||||
'langbot.pkg.core.app': MagicMock(),
|
||||
'langbot_plugin.api.entities.builtin.resource.tool': MagicMock(),
|
||||
'langbot_plugin.api.entities.builtin.pipeline.query': MagicMock(),
|
||||
'langbot_plugin.api.entities.builtin.provider.message': MagicMock(),
|
||||
'langbot.pkg.provider.modelmgr.errors': MagicMock(),
|
||||
}
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.provider.modelmgr.requesters.chatcmpl import OpenAIChatCompletions
|
||||
|
||||
mock_app = MagicMock()
|
||||
requester = OpenAIChatCompletions(mock_app, {})
|
||||
return requester
|
||||
|
||||
def test_mask_api_key_full(self):
|
||||
"""Mask a full API key."""
|
||||
requester = self._create_requester_with_mocks()
|
||||
|
||||
result = requester._mask_api_key('sk-1234567890abcdef')
|
||||
assert result == 'sk-1...cdef'
|
||||
|
||||
def test_mask_api_key_short(self):
|
||||
"""Mask a short API key (<=8 chars)."""
|
||||
requester = self._create_requester_with_mocks()
|
||||
|
||||
result = requester._mask_api_key('short')
|
||||
assert result == '****'
|
||||
|
||||
def test_mask_api_key_empty(self):
|
||||
"""Empty API key returns empty string."""
|
||||
requester = self._create_requester_with_mocks()
|
||||
|
||||
result = requester._mask_api_key('')
|
||||
assert result == ''
|
||||
|
||||
def test_mask_api_key_none(self):
|
||||
"""None API key returns empty string."""
|
||||
requester = self._create_requester_with_mocks()
|
||||
|
||||
result = requester._mask_api_key(None)
|
||||
assert result == ''
|
||||
|
||||
def test_mask_api_key_exact_8_chars(self):
|
||||
"""API key with exactly 8 chars is masked as **** (<=8 threshold)."""
|
||||
requester = self._create_requester_with_mocks()
|
||||
|
||||
result = requester._mask_api_key('12345678')
|
||||
assert result == '****' # <= 8 chars gets masked
|
||||
|
||||
|
||||
class TestInferModelType:
|
||||
"""Tests for _infer_model_type method."""
|
||||
|
||||
def _create_requester_with_mocks(self):
|
||||
mocks = {
|
||||
'langbot.pkg.core.app': MagicMock(),
|
||||
'langbot_plugin.api.entities.builtin.resource.tool': MagicMock(),
|
||||
'langbot_plugin.api.entities.builtin.pipeline.query': MagicMock(),
|
||||
'langbot_plugin.api.entities.builtin.provider.message': MagicMock(),
|
||||
'langbot.pkg.provider.modelmgr.errors': MagicMock(),
|
||||
}
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.provider.modelmgr.requesters.chatcmpl import OpenAIChatCompletions
|
||||
|
||||
mock_app = MagicMock()
|
||||
requester = OpenAIChatCompletions(mock_app, {})
|
||||
return requester
|
||||
|
||||
def test_infer_embedding_from_name(self):
|
||||
"""Infer embedding type from model name."""
|
||||
requester = self._create_requester_with_mocks()
|
||||
|
||||
assert requester._infer_model_type('text-embedding-ada-002') == 'embedding'
|
||||
assert requester._infer_model_type('bge-large-en') == 'embedding'
|
||||
assert requester._infer_model_type('e5-base') == 'embedding'
|
||||
assert requester._infer_model_type('m3e-base') == 'embedding'
|
||||
|
||||
def test_infer_llm_from_name(self):
|
||||
"""Infer LLM type from model name."""
|
||||
requester = self._create_requester_with_mocks()
|
||||
|
||||
assert requester._infer_model_type('gpt-4') == 'llm'
|
||||
assert requester._infer_model_type('claude-3-opus') == 'llm'
|
||||
assert requester._infer_model_type('llama-2-70b') == 'llm'
|
||||
|
||||
def test_infer_model_type_none_id(self):
|
||||
"""Handle None model_id."""
|
||||
requester = self._create_requester_with_mocks()
|
||||
|
||||
result = requester._infer_model_type(None)
|
||||
assert result == 'llm' # Default
|
||||
|
||||
def test_infer_model_type_empty_id(self):
|
||||
"""Handle empty model_id."""
|
||||
requester = self._create_requester_with_mocks()
|
||||
|
||||
result = requester._infer_model_type('')
|
||||
assert result == 'llm' # Default
|
||||
|
||||
|
||||
class TestNormalizeModalities:
|
||||
"""Tests for _normalize_modalities method."""
|
||||
|
||||
def _create_requester_with_mocks(self):
|
||||
mocks = {
|
||||
'langbot.pkg.core.app': MagicMock(),
|
||||
'langbot_plugin.api.entities.builtin.resource.tool': MagicMock(),
|
||||
'langbot_plugin.api.entities.builtin.pipeline.query': MagicMock(),
|
||||
'langbot_plugin.api.entities.builtin.provider.message': MagicMock(),
|
||||
'langbot.pkg.provider.modelmgr.errors': MagicMock(),
|
||||
}
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.provider.modelmgr.requesters.chatcmpl import OpenAIChatCompletions
|
||||
|
||||
mock_app = MagicMock()
|
||||
requester = OpenAIChatCompletions(mock_app, {})
|
||||
return requester
|
||||
|
||||
def test_normalize_string_modality(self):
|
||||
"""Normalize single string modality."""
|
||||
requester = self._create_requester_with_mocks()
|
||||
|
||||
result = requester._normalize_modalities('text,image')
|
||||
assert result == ['text', 'image']
|
||||
|
||||
def test_normalize_list_modalities(self):
|
||||
"""Normalize list of modalities."""
|
||||
requester = self._create_requester_with_mocks()
|
||||
|
||||
result = requester._normalize_modalities(['text', 'image', 'audio'])
|
||||
assert result == ['text', 'image', 'audio']
|
||||
|
||||
def test_normalize_dict_modalities(self):
|
||||
"""Normalize dict with nested modalities."""
|
||||
requester = self._create_requester_with_mocks()
|
||||
|
||||
result = requester._normalize_modalities({'input': ['text'], 'output': ['text', 'image']})
|
||||
assert result == ['text', 'image']
|
||||
|
||||
def test_normalize_none(self):
|
||||
"""Handle None input."""
|
||||
requester = self._create_requester_with_mocks()
|
||||
|
||||
result = requester._normalize_modalities(None)
|
||||
assert result == []
|
||||
|
||||
def test_normalize_arrow_separator(self):
|
||||
"""Handle arrow separator in modality string."""
|
||||
requester = self._create_requester_with_mocks()
|
||||
|
||||
result = requester._normalize_modalities('text->image')
|
||||
assert result == ['text', 'image']
|
||||
|
||||
|
||||
class TestParseRerankResponse:
|
||||
"""Tests for _parse_rerank_response static method."""
|
||||
|
||||
def test_parse_cohere_jina_format(self):
|
||||
"""Parse Cohere/Jina/SiliconFlow format."""
|
||||
from langbot.pkg.provider.modelmgr.requesters.chatcmpl import OpenAIChatCompletions
|
||||
|
||||
data = {
|
||||
'results': [
|
||||
{'index': 0, 'relevance_score': 0.95},
|
||||
{'index': 1, 'relevance_score': 0.80},
|
||||
]
|
||||
}
|
||||
|
||||
result = OpenAIChatCompletions._parse_rerank_response(data)
|
||||
assert result == [
|
||||
{'index': 0, 'relevance_score': 0.95},
|
||||
{'index': 1, 'relevance_score': 0.80},
|
||||
]
|
||||
|
||||
def test_parse_voyage_format(self):
|
||||
"""Parse Voyage AI format."""
|
||||
from langbot.pkg.provider.modelmgr.requesters.chatcmpl import OpenAIChatCompletions
|
||||
|
||||
data = {
|
||||
'data': [
|
||||
{'index': 0, 'relevance_score': 0.90},
|
||||
{'index': 2, 'relevance_score': 0.75},
|
||||
]
|
||||
}
|
||||
|
||||
result = OpenAIChatCompletions._parse_rerank_response(data)
|
||||
assert result == [
|
||||
{'index': 0, 'relevance_score': 0.90},
|
||||
{'index': 2, 'relevance_score': 0.75},
|
||||
]
|
||||
|
||||
def test_parse_dashscope_format(self):
|
||||
"""Parse DashScope format."""
|
||||
from langbot.pkg.provider.modelmgr.requesters.chatcmpl import OpenAIChatCompletions
|
||||
|
||||
data = {
|
||||
'output': {
|
||||
'results': [
|
||||
{'index': 0, 'relevance_score': 0.85},
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
result = OpenAIChatCompletions._parse_rerank_response(data)
|
||||
assert result == [{'index': 0, 'relevance_score': 0.85}]
|
||||
|
||||
def test_parse_unknown_format(self):
|
||||
"""Handle unknown format returns empty list."""
|
||||
from langbot.pkg.provider.modelmgr.requesters.chatcmpl import OpenAIChatCompletions
|
||||
|
||||
data = {'unknown_key': 'value'}
|
||||
|
||||
result = OpenAIChatCompletions._parse_rerank_response(data)
|
||||
assert result == []
|
||||
|
||||
def test_parse_empty_results(self):
|
||||
"""Handle empty results."""
|
||||
from langbot.pkg.provider.modelmgr.requesters.chatcmpl import OpenAIChatCompletions
|
||||
|
||||
data = {'results': []}
|
||||
|
||||
result = OpenAIChatCompletions._parse_rerank_response(data)
|
||||
assert result == []
|
||||
|
||||
|
||||
class TestExtractScanMetadata:
|
||||
"""Tests for _extract_scan_metadata method."""
|
||||
|
||||
def _create_requester_with_mocks(self):
|
||||
mocks = {
|
||||
'langbot.pkg.core.app': MagicMock(),
|
||||
'langbot_plugin.api.entities.builtin.resource.tool': MagicMock(),
|
||||
'langbot_plugin.api.entities.builtin.pipeline.query': MagicMock(),
|
||||
'langbot_plugin.api.entities.builtin.provider.message': MagicMock(),
|
||||
'langbot.pkg.provider.modelmgr.errors': MagicMock(),
|
||||
}
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.provider.modelmgr.requesters.chatcmpl import OpenAIChatCompletions
|
||||
|
||||
mock_app = MagicMock()
|
||||
requester = OpenAIChatCompletions(mock_app, {})
|
||||
return requester
|
||||
|
||||
def test_extract_basic_metadata(self):
|
||||
"""Extract basic model metadata."""
|
||||
requester = self._create_requester_with_mocks()
|
||||
|
||||
item = {
|
||||
'id': 'gpt-4',
|
||||
'name': 'GPT-4 Turbo',
|
||||
'description': 'Most capable GPT-4 model',
|
||||
'context_length': 128000,
|
||||
'owned_by': 'openai',
|
||||
}
|
||||
|
||||
result = requester._extract_scan_metadata(item, 'gpt-4')
|
||||
|
||||
assert result['display_name'] == 'GPT-4 Turbo'
|
||||
assert result['description'] == 'Most capable GPT-4 model'
|
||||
assert result['context_length'] == 128000
|
||||
assert result['owned_by'] == 'openai'
|
||||
|
||||
def test_extract_metadata_missing_fields(self):
|
||||
"""Handle missing metadata fields."""
|
||||
requester = self._create_requester_with_mocks()
|
||||
|
||||
item = {'id': 'unknown-model'}
|
||||
|
||||
result = requester._extract_scan_metadata(item, 'unknown-model')
|
||||
|
||||
assert result['display_name'] is None
|
||||
assert result['description'] is None
|
||||
assert result['context_length'] is None
|
||||
assert result['owned_by'] is None
|
||||
|
||||
def test_extract_metadata_top_provider_context(self):
|
||||
"""Extract context_length from top_provider."""
|
||||
requester = self._create_requester_with_mocks()
|
||||
|
||||
item = {
|
||||
'id': 'model',
|
||||
'top_provider': {
|
||||
'context_length': 4096,
|
||||
},
|
||||
}
|
||||
|
||||
result = requester._extract_scan_metadata(item, 'model')
|
||||
|
||||
assert result['context_length'] == 4096
|
||||
|
||||
def test_extract_metadata_empty_strings(self):
|
||||
"""Handle empty string values."""
|
||||
requester = self._create_requester_with_mocks()
|
||||
|
||||
item = {
|
||||
'id': 'model',
|
||||
'name': '', # Empty name
|
||||
'description': ' ', # Whitespace only
|
||||
'owned_by': '',
|
||||
}
|
||||
|
||||
result = requester._extract_scan_metadata(item, 'model')
|
||||
|
||||
assert result['display_name'] is None
|
||||
assert result['description'] is None
|
||||
assert result['owned_by'] is None
|
||||
|
||||
def test_extract_metadata_name_matches_id(self):
|
||||
"""When name equals id, display_name is None."""
|
||||
requester = self._create_requester_with_mocks()
|
||||
|
||||
item = {
|
||||
'id': 'gpt-4',
|
||||
'name': 'gpt-4', # Same as id
|
||||
}
|
||||
|
||||
result = requester._extract_scan_metadata(item, 'gpt-4')
|
||||
|
||||
assert result['display_name'] is None
|
||||
@@ -1,264 +0,0 @@
|
||||
"""Tests for OllamaChatCompletions requester.
|
||||
|
||||
Tests model inference, payload construction, and error handling.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.provider.modelmgr.errors import RequesterError
|
||||
|
||||
|
||||
class TestOllamaRequesterConfig:
|
||||
"""Tests for default config."""
|
||||
|
||||
def test_default_config_values(self):
|
||||
"""Check default_config."""
|
||||
from langbot.pkg.provider.modelmgr.requesters.ollamachat import OllamaChatCompletions
|
||||
|
||||
assert OllamaChatCompletions.default_config['base_url'] == 'http://127.0.0.1:11434'
|
||||
assert OllamaChatCompletions.default_config['timeout'] == 120
|
||||
|
||||
def test_config_override(self):
|
||||
"""Config can override defaults."""
|
||||
from langbot.pkg.provider.modelmgr.requesters.ollamachat import OllamaChatCompletions
|
||||
|
||||
mock_app = MagicMock()
|
||||
req = OllamaChatCompletions(mock_app, {
|
||||
'base_url': 'http://custom.ollama:11434',
|
||||
'timeout': 300,
|
||||
})
|
||||
|
||||
assert req.requester_cfg['base_url'] == 'http://custom.ollama:11434'
|
||||
assert req.requester_cfg['timeout'] == 300
|
||||
|
||||
|
||||
class TestOllamaInferModelType:
|
||||
"""Tests for _infer_model_type pure function."""
|
||||
|
||||
@pytest.fixture
|
||||
def requester(self):
|
||||
from langbot.pkg.provider.modelmgr.requesters.ollamachat import OllamaChatCompletions
|
||||
|
||||
return OllamaChatCompletions(MagicMock(), {})
|
||||
|
||||
def test_infer_embedding_from_name(self, requester):
|
||||
"""Embedding keywords return 'embedding'."""
|
||||
assert requester._infer_model_type('nomic-embed-text') == 'embedding'
|
||||
assert requester._infer_model_type('bge-large') == 'embedding'
|
||||
assert requester._infer_model_type('text-embedding') == 'embedding'
|
||||
|
||||
def test_infer_llm_from_name(self, requester):
|
||||
"""Non-embedding keywords return 'llm'."""
|
||||
assert requester._infer_model_type('llama2') == 'llm'
|
||||
assert requester._infer_model_type('mistral') == 'llm'
|
||||
assert requester._infer_model_type('codellama') == 'llm'
|
||||
|
||||
def test_infer_model_type_none(self, requester):
|
||||
"""None model_id returns 'llm'."""
|
||||
assert requester._infer_model_type(None) == 'llm'
|
||||
|
||||
def test_infer_model_type_empty(self, requester):
|
||||
"""Empty model_id returns 'llm'."""
|
||||
assert requester._infer_model_type('') == 'llm'
|
||||
|
||||
|
||||
class TestOllamaInferModelAbilities:
|
||||
"""Tests for _infer_model_abilities pure function."""
|
||||
|
||||
@pytest.fixture
|
||||
def requester(self):
|
||||
from langbot.pkg.provider.modelmgr.requesters.ollamachat import OllamaChatCompletions
|
||||
|
||||
return OllamaChatCompletions(MagicMock(), {})
|
||||
|
||||
def test_infer_vision_ability(self, requester):
|
||||
"""Vision keywords add 'vision' ability."""
|
||||
item = {
|
||||
'details': {
|
||||
'family': 'llava',
|
||||
}
|
||||
}
|
||||
|
||||
abilities = requester._infer_model_abilities(item, 'llava-v1.5')
|
||||
assert 'vision' in abilities
|
||||
|
||||
def test_infer_vision_from_model_id(self, requester):
|
||||
"""Vision keywords in model_id add 'vision' ability."""
|
||||
item = {}
|
||||
abilities = requester._infer_model_abilities(item, 'llava-7b')
|
||||
assert 'vision' in abilities
|
||||
|
||||
def test_infer_func_call_ability(self, requester):
|
||||
"""Tool/function keywords add 'func_call' ability."""
|
||||
item = {
|
||||
'details': {
|
||||
'families': ['tools'],
|
||||
}
|
||||
}
|
||||
|
||||
abilities = requester._infer_model_abilities(item, 'model')
|
||||
assert 'func_call' in abilities
|
||||
|
||||
def test_infer_no_abilities(self, requester):
|
||||
"""No matching keywords returns empty abilities."""
|
||||
item = {
|
||||
'details': {
|
||||
'family': 'llama',
|
||||
}
|
||||
}
|
||||
|
||||
abilities = requester._infer_model_abilities(item, 'llama-2')
|
||||
assert len(abilities) == 0
|
||||
|
||||
def test_infer_multiple_abilities(self, requester):
|
||||
"""Multiple keywords can add multiple abilities."""
|
||||
item = {
|
||||
'details': {
|
||||
'family': 'vision',
|
||||
'families': ['tools'],
|
||||
}
|
||||
}
|
||||
|
||||
abilities = requester._infer_model_abilities(item, 'vision-tool-model')
|
||||
assert 'vision' in abilities
|
||||
assert 'func_call' in abilities
|
||||
|
||||
|
||||
class TestOllamaMakeMessage:
|
||||
"""Tests for _make_msg response parsing."""
|
||||
|
||||
@pytest.fixture
|
||||
def requester(self):
|
||||
from langbot.pkg.provider.modelmgr.requesters.ollamachat import OllamaChatCompletions
|
||||
|
||||
return OllamaChatCompletions(MagicMock(), {})
|
||||
|
||||
def _create_ollama_response(self, content, tool_calls=None):
|
||||
"""Helper to create mock ollama response."""
|
||||
import ollama
|
||||
|
||||
mock_response = MagicMock(spec=ollama.ChatResponse)
|
||||
mock_message = MagicMock(spec=ollama.Message)
|
||||
mock_message.content = content
|
||||
mock_message.tool_calls = tool_calls
|
||||
mock_response.message = mock_message
|
||||
|
||||
return mock_response
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_make_msg_text_content(self, requester):
|
||||
"""Text content is extracted."""
|
||||
mock_response = self._create_ollama_response('Hello world')
|
||||
|
||||
result = await requester._make_msg(mock_response)
|
||||
|
||||
assert result.content == 'Hello world'
|
||||
assert result.role == 'assistant'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_make_msg_with_tool_calls(self, requester):
|
||||
"""Tool calls are parsed."""
|
||||
mock_tool_call = MagicMock()
|
||||
mock_tool_call.function = MagicMock()
|
||||
mock_tool_call.function.name = 'get_weather'
|
||||
mock_tool_call.function.arguments = {'location': 'Beijing'}
|
||||
|
||||
mock_response = self._create_ollama_response('', tool_calls=[mock_tool_call])
|
||||
|
||||
result = await requester._make_msg(mock_response)
|
||||
|
||||
assert result.tool_calls is not None
|
||||
assert len(result.tool_calls) == 1
|
||||
assert result.tool_calls[0].function.name == 'get_weather'
|
||||
# Arguments should be JSON string
|
||||
assert isinstance(result.tool_calls[0].function.arguments, str)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_make_msg_empty_message_raises(self, requester):
|
||||
"""Empty message raises ValueError."""
|
||||
mock_response = MagicMock()
|
||||
mock_response.message = None
|
||||
|
||||
with pytest.raises(ValueError, match='message'):
|
||||
await requester._make_msg(mock_response)
|
||||
|
||||
|
||||
class TestOllamaErrorHandling:
|
||||
"""Tests for error handling branches."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_app(self):
|
||||
app = MagicMock()
|
||||
app.tool_mgr = MagicMock()
|
||||
app.tool_mgr.generate_tools_for_openai = AsyncMock(return_value=[])
|
||||
return app
|
||||
|
||||
@pytest.fixture
|
||||
def requester_with_mocked_client(self, mock_app):
|
||||
from langbot.pkg.provider.modelmgr.requesters.ollamachat import OllamaChatCompletions
|
||||
|
||||
req = OllamaChatCompletions(mock_app, {})
|
||||
req.client = MagicMock()
|
||||
req.client.chat = AsyncMock()
|
||||
|
||||
return req
|
||||
|
||||
@pytest.fixture
|
||||
def mock_model(self):
|
||||
model = MagicMock()
|
||||
model.model_entity = MagicMock()
|
||||
model.model_entity.name = 'llama2'
|
||||
model.provider = MagicMock()
|
||||
model.provider.token_mgr = MagicMock()
|
||||
model.provider.token_mgr.get_token = MagicMock(return_value='')
|
||||
return model
|
||||
|
||||
@pytest.fixture
|
||||
def mock_message(self):
|
||||
msg = MagicMock()
|
||||
msg.role = 'user'
|
||||
msg.content = 'test'
|
||||
msg.dict = MagicMock(return_value={'role': 'user', 'content': 'test'})
|
||||
return msg
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_timeout_error(self, requester_with_mocked_client, mock_model, mock_message):
|
||||
"""TimeoutError is converted to RequesterError."""
|
||||
requester_with_mocked_client.client.chat = AsyncMock(side_effect=asyncio.TimeoutError())
|
||||
|
||||
with pytest.raises(RequesterError) as exc:
|
||||
await requester_with_mocked_client.invoke_llm(
|
||||
query=None,
|
||||
model=mock_model,
|
||||
messages=[mock_message],
|
||||
)
|
||||
|
||||
assert '超时' in str(exc.value)
|
||||
|
||||
|
||||
class TestOllamaScanModels:
|
||||
"""Tests for scan_models method."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_app(self):
|
||||
return MagicMock()
|
||||
|
||||
@pytest.fixture
|
||||
def requester(self, mock_app):
|
||||
from langbot.pkg.provider.modelmgr.requesters.ollamachat import OllamaChatCompletions
|
||||
|
||||
req = OllamaChatCompletions(mock_app, {
|
||||
'base_url': 'http://127.0.0.1:11434',
|
||||
'timeout': 120,
|
||||
})
|
||||
return req
|
||||
|
||||
def test_requester_name_constant(self):
|
||||
"""REQUESTER_NAME constant exists."""
|
||||
from langbot.pkg.provider.modelmgr.requesters.ollamachat import REQUESTER_NAME
|
||||
|
||||
assert REQUESTER_NAME == 'ollama-chat'
|
||||
@@ -0,0 +1,93 @@
|
||||
"""Unit tests for LiteLLMRequester._convert_messages.
|
||||
|
||||
Focus: the content-part normalization that (a) converts image_base64 parts to
|
||||
the OpenAI image_url shape and (b) drops non-image file parts (file_base64 /
|
||||
file_url) which OpenAI-compatible chat models reject. The latter is essential
|
||||
for Voice/File attachments — including ones replayed from conversation history —
|
||||
since the agent consumes their bytes via the sandbox, not the model payload.
|
||||
"""
|
||||
|
||||
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
||||
|
||||
from langbot.pkg.provider.modelmgr.requesters.litellmchat import LiteLLMRequester
|
||||
|
||||
|
||||
def _make_requester() -> LiteLLMRequester:
|
||||
# _convert_messages does not touch instance config, so bypass __init__.
|
||||
return LiteLLMRequester.__new__(LiteLLMRequester)
|
||||
|
||||
|
||||
def test_convert_messages_drops_file_base64_part():
|
||||
req = _make_requester()
|
||||
msg = provider_message.Message(
|
||||
role='user',
|
||||
content=[
|
||||
provider_message.ContentElement.from_text('analyze this audio'),
|
||||
provider_message.ContentElement.from_file_base64('data:audio/wav;base64,AAAA', 'voice.wav'),
|
||||
],
|
||||
)
|
||||
out = req._convert_messages([msg])
|
||||
parts = out[0]['content']
|
||||
types = [p.get('type') for p in parts]
|
||||
assert 'file_base64' not in types
|
||||
assert types == ['text']
|
||||
assert parts[0]['text'] == 'analyze this audio'
|
||||
|
||||
|
||||
def test_convert_messages_drops_file_url_part():
|
||||
req = _make_requester()
|
||||
msg = provider_message.Message(
|
||||
role='user',
|
||||
content=[
|
||||
provider_message.ContentElement.from_text('here is a doc'),
|
||||
provider_message.ContentElement.from_file_url('http://example.com/report.xlsx', 'report.xlsx'),
|
||||
],
|
||||
)
|
||||
out = req._convert_messages([msg])
|
||||
types = [p.get('type') for p in out[0]['content']]
|
||||
assert types == ['text']
|
||||
|
||||
|
||||
def test_convert_messages_keeps_image_and_converts_to_image_url():
|
||||
req = _make_requester()
|
||||
msg = provider_message.Message(
|
||||
role='user',
|
||||
content=[
|
||||
provider_message.ContentElement.from_text('look'),
|
||||
provider_message.ContentElement.from_image_base64('data:image/png;base64,AAAA'),
|
||||
],
|
||||
)
|
||||
out = req._convert_messages([msg])
|
||||
parts = out[0]['content']
|
||||
types = [p.get('type') for p in parts]
|
||||
# image is preserved and reshaped to the OpenAI image_url form
|
||||
assert types == ['text', 'image_url']
|
||||
img_part = parts[1]
|
||||
assert img_part['image_url'] == {'url': 'data:image/png;base64,AAAA'}
|
||||
assert 'image_base64' not in img_part
|
||||
|
||||
|
||||
def test_convert_messages_mixed_history_strips_only_files():
|
||||
req = _make_requester()
|
||||
# Simulate replayed history: an old voice turn + a current text turn.
|
||||
history_voice = provider_message.Message(
|
||||
role='user',
|
||||
content=[
|
||||
provider_message.ContentElement.from_text('old audio turn'),
|
||||
provider_message.ContentElement.from_file_base64('data:audio/wav;base64,BBBB', 'voice.wav'),
|
||||
],
|
||||
)
|
||||
current = provider_message.Message(
|
||||
role='user',
|
||||
content=[provider_message.ContentElement.from_text('now do the csv')],
|
||||
)
|
||||
out = req._convert_messages([history_voice, current])
|
||||
assert [p.get('type') for p in out[0]['content']] == ['text']
|
||||
assert [p.get('type') for p in out[1]['content']] == ['text']
|
||||
|
||||
|
||||
def test_convert_messages_plain_string_content_untouched():
|
||||
req = _make_requester()
|
||||
msg = provider_message.Message(role='user', content='just text')
|
||||
out = req._convert_messages([msg])
|
||||
assert out[0]['content'] == 'just text'
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,146 @@
|
||||
"""Unit tests for LocalAgentRunner._inject_inbound_attachments.
|
||||
|
||||
Covers the user -> sandbox attachment path added for the Box attachment
|
||||
round-trip:
|
||||
|
||||
* materialized descriptors are stashed on the query and described to the model
|
||||
via an appended text note (in-sandbox paths + outbox convention);
|
||||
* non-image file parts (file_base64 / file_url) are stripped from the user
|
||||
message content because OpenAI-compatible chat models reject them, while
|
||||
image and text parts are kept for vision models;
|
||||
* the helper is a no-op when the box service is unavailable or yields nothing,
|
||||
and never raises into the chat turn on materialization failure.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
|
||||
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
||||
|
||||
from langbot.pkg.provider.runners.localagent import LocalAgentRunner
|
||||
|
||||
|
||||
def _make_runner(box_service) -> LocalAgentRunner:
|
||||
runner = LocalAgentRunner.__new__(LocalAgentRunner)
|
||||
runner.ap = SimpleNamespace(logger=Mock(), box_service=box_service)
|
||||
return runner
|
||||
|
||||
|
||||
def _make_query():
|
||||
return SimpleNamespace(variables={}, query_id='q-123')
|
||||
|
||||
|
||||
def _box_service(attachments):
|
||||
svc = SimpleNamespace(
|
||||
available=True,
|
||||
OUTBOX_MOUNT_DIR='/outbox',
|
||||
materialize_inbound_attachments=AsyncMock(return_value=attachments),
|
||||
)
|
||||
return svc
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inject_strips_file_parts_and_appends_note():
|
||||
box = _box_service([{'type': 'Voice', 'path': '/inbox/q-123/voice.wav', 'size': 176000}])
|
||||
runner = _make_runner(box)
|
||||
query = _make_query()
|
||||
user_message = provider_message.Message(
|
||||
role='user',
|
||||
content=[
|
||||
provider_message.ContentElement.from_text('transcribe this'),
|
||||
provider_message.ContentElement.from_file_base64('data:audio/wav;base64,AAAA', 'voice.wav'),
|
||||
],
|
||||
)
|
||||
|
||||
await runner._inject_inbound_attachments(query, user_message)
|
||||
|
||||
types = [getattr(ce, 'type', None) for ce in user_message.content]
|
||||
# file_base64 dropped; text kept; sandbox-path note appended as text
|
||||
assert 'file_base64' not in types
|
||||
assert types.count('text') == 2
|
||||
note = user_message.content[-1].text
|
||||
assert '/inbox/q-123/voice.wav' in note
|
||||
assert '/outbox/q-123' in note
|
||||
# descriptors stashed for downstream stages
|
||||
assert query.variables['_sandbox_inbound_attachments'] == box.materialize_inbound_attachments.return_value
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inject_keeps_image_parts():
|
||||
box = _box_service([{'type': 'Image', 'path': '/inbox/q-123/pic.png', 'size': 1234}])
|
||||
runner = _make_runner(box)
|
||||
query = _make_query()
|
||||
user_message = provider_message.Message(
|
||||
role='user',
|
||||
content=[
|
||||
provider_message.ContentElement.from_text('what is this'),
|
||||
provider_message.ContentElement.from_image_base64('data:image/png;base64,iVBORw0K'),
|
||||
],
|
||||
)
|
||||
|
||||
await runner._inject_inbound_attachments(query, user_message)
|
||||
|
||||
types = [getattr(ce, 'type', None) for ce in user_message.content]
|
||||
assert 'image_base64' in types # vision part preserved
|
||||
assert types[-1] == 'text' # note appended last
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inject_promotes_string_content_to_list_with_note():
|
||||
box = _box_service([{'type': 'File', 'path': '/inbox/q-123/data.csv', 'size': 42}])
|
||||
runner = _make_runner(box)
|
||||
query = _make_query()
|
||||
user_message = provider_message.Message(role='user', content='clean this csv')
|
||||
|
||||
await runner._inject_inbound_attachments(query, user_message)
|
||||
|
||||
assert isinstance(user_message.content, list)
|
||||
assert [getattr(ce, 'type', None) for ce in user_message.content] == ['text', 'text']
|
||||
assert user_message.content[0].text == 'clean this csv'
|
||||
assert '/inbox/q-123/data.csv' in user_message.content[1].text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inject_noop_without_box_service():
|
||||
runner = _make_runner(box_service=None)
|
||||
query = _make_query()
|
||||
user_message = provider_message.Message(role='user', content='hello')
|
||||
|
||||
await runner._inject_inbound_attachments(query, user_message)
|
||||
|
||||
assert user_message.content == 'hello'
|
||||
assert '_sandbox_inbound_attachments' not in query.variables
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inject_noop_when_no_attachments():
|
||||
box = _box_service([])
|
||||
runner = _make_runner(box)
|
||||
query = _make_query()
|
||||
user_message = provider_message.Message(role='user', content='hello')
|
||||
|
||||
await runner._inject_inbound_attachments(query, user_message)
|
||||
|
||||
assert user_message.content == 'hello'
|
||||
assert '_sandbox_inbound_attachments' not in query.variables
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inject_swallows_materialization_error():
|
||||
box = SimpleNamespace(
|
||||
available=True,
|
||||
OUTBOX_MOUNT_DIR='/outbox',
|
||||
materialize_inbound_attachments=AsyncMock(side_effect=RuntimeError('disk full')),
|
||||
)
|
||||
runner = _make_runner(box)
|
||||
query = _make_query()
|
||||
user_message = provider_message.Message(role='user', content='hello')
|
||||
|
||||
# must not raise
|
||||
await runner._inject_inbound_attachments(query, user_message)
|
||||
assert user_message.content == 'hello'
|
||||
runner.ap.logger.warning.assert_called_once()
|
||||
@@ -0,0 +1,281 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
|
||||
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.provider.runners.localagent import LocalAgentRunner, _StreamAccumulator
|
||||
|
||||
|
||||
class RecordingProvider:
|
||||
def __init__(self):
|
||||
self.requests: list[dict] = []
|
||||
|
||||
async def invoke_llm(self, query, model, messages, funcs, extra_args=None, remove_think=None):
|
||||
self.requests.append(
|
||||
{
|
||||
'messages': list(messages),
|
||||
'funcs': list(funcs),
|
||||
'remove_think': remove_think,
|
||||
}
|
||||
)
|
||||
|
||||
if len(self.requests) == 1:
|
||||
return provider_message.Message(
|
||||
role='assistant',
|
||||
content='Let me calculate that exactly.',
|
||||
tool_calls=[
|
||||
provider_message.ToolCall(
|
||||
id='call-1',
|
||||
type='function',
|
||||
function=provider_message.FunctionCall(
|
||||
name='exec',
|
||||
arguments=json.dumps(
|
||||
{'command': ("python - <<'PY'\nnums = [1, 2, 3, 4]\nprint(sum(nums) / len(nums))\nPY")}
|
||||
),
|
||||
),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
tool_result = json.loads(messages[-1].content)
|
||||
return provider_message.Message(
|
||||
role='assistant',
|
||||
content=f'The average is {tool_result["stdout"]}.',
|
||||
)
|
||||
|
||||
|
||||
class RecordingStreamProvider:
|
||||
def __init__(self):
|
||||
self.stream_requests: list[dict] = []
|
||||
|
||||
def invoke_llm_stream(self, query, model, messages, funcs, extra_args=None, remove_think=None):
|
||||
self.stream_requests.append(
|
||||
{
|
||||
'messages': list(messages),
|
||||
'funcs': list(funcs),
|
||||
'remove_think': remove_think,
|
||||
}
|
||||
)
|
||||
|
||||
async def _stream():
|
||||
if len(self.stream_requests) == 1:
|
||||
yield provider_message.MessageChunk(
|
||||
role='assistant',
|
||||
tool_calls=[
|
||||
provider_message.ToolCall(
|
||||
id='call-1',
|
||||
type='function',
|
||||
function=provider_message.FunctionCall(
|
||||
name='exec',
|
||||
arguments=json.dumps({'command': "python -c 'print(1)'"}),
|
||||
),
|
||||
)
|
||||
],
|
||||
is_final=True,
|
||||
)
|
||||
return
|
||||
|
||||
yield provider_message.MessageChunk(
|
||||
role='assistant',
|
||||
content='Tool execution failed.',
|
||||
is_final=True,
|
||||
)
|
||||
|
||||
return _stream()
|
||||
|
||||
|
||||
def make_query() -> pipeline_query.Query:
|
||||
adapter = AsyncMock()
|
||||
adapter.is_stream_output_supported = AsyncMock(return_value=False)
|
||||
|
||||
return pipeline_query.Query.model_construct(
|
||||
query_id='avg-query',
|
||||
launcher_type=provider_session.LauncherTypes.PERSON,
|
||||
launcher_id=12345,
|
||||
sender_id=12345,
|
||||
message_chain=[],
|
||||
message_event=None,
|
||||
adapter=adapter,
|
||||
pipeline_uuid='pipeline-uuid',
|
||||
bot_uuid='bot-uuid',
|
||||
pipeline_config={
|
||||
'ai': {
|
||||
'runner': {'runner': 'local-agent'},
|
||||
'local-agent': {'model': {'primary': 'test-model-uuid', 'fallbacks': []}, 'prompt': 'test-prompt'},
|
||||
},
|
||||
'output': {'misc': {'remove-think': False}},
|
||||
},
|
||||
prompt=SimpleNamespace(messages=[]),
|
||||
messages=[],
|
||||
user_message=provider_message.Message(
|
||||
role='user',
|
||||
content='Please calculate the average of 1, 2, 3, and 4.',
|
||||
),
|
||||
use_funcs=[SimpleNamespace(name='exec')],
|
||||
use_llm_model_uuid='test-model-uuid',
|
||||
variables={},
|
||||
)
|
||||
|
||||
|
||||
def test_stream_accumulator_merges_fragmented_tool_call_arguments():
|
||||
accumulator = _StreamAccumulator(msg_sequence=1)
|
||||
|
||||
assert (
|
||||
accumulator.add(
|
||||
provider_message.MessageChunk(
|
||||
role='assistant',
|
||||
tool_calls=[
|
||||
provider_message.ToolCall(
|
||||
id='call-1',
|
||||
type='function',
|
||||
function=provider_message.FunctionCall(name='exec', arguments='{"command":'),
|
||||
)
|
||||
],
|
||||
)
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
emitted = accumulator.add(
|
||||
provider_message.MessageChunk(
|
||||
role='assistant',
|
||||
tool_calls=[
|
||||
provider_message.ToolCall(
|
||||
id='call-1',
|
||||
type='function',
|
||||
function=provider_message.FunctionCall(name='exec', arguments='"pwd"}'),
|
||||
)
|
||||
],
|
||||
is_final=True,
|
||||
)
|
||||
)
|
||||
|
||||
assert emitted is not None
|
||||
final_msg = accumulator.final_message()
|
||||
assert final_msg.tool_calls[0].function.name == 'exec'
|
||||
assert final_msg.tool_calls[0].function.arguments == '{"command":"pwd"}'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_localagent_uses_exec_for_exact_calculation():
|
||||
provider = RecordingProvider()
|
||||
model = SimpleNamespace(
|
||||
provider=provider,
|
||||
model_entity=SimpleNamespace(
|
||||
uuid='test-model-uuid',
|
||||
name='test-model',
|
||||
abilities=['func_call'],
|
||||
extra_args={},
|
||||
),
|
||||
)
|
||||
|
||||
tool_manager = SimpleNamespace(
|
||||
execute_func_call=AsyncMock(
|
||||
return_value={
|
||||
'session_id': 'avg-query',
|
||||
'backend': 'podman',
|
||||
'status': 'completed',
|
||||
'ok': True,
|
||||
'exit_code': 0,
|
||||
'stdout': '2.5',
|
||||
'stderr': '',
|
||||
'duration_ms': 18,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
app = SimpleNamespace(
|
||||
logger=Mock(),
|
||||
model_mgr=SimpleNamespace(get_model_by_uuid=AsyncMock(return_value=model)),
|
||||
tool_mgr=tool_manager,
|
||||
rag_mgr=SimpleNamespace(),
|
||||
box_service=SimpleNamespace(
|
||||
get_system_guidance=Mock(
|
||||
return_value=(
|
||||
'When the exec tool is available, use it for exact calculations, statistics, '
|
||||
'structured data parsing, and code execution instead of estimating mentally. '
|
||||
'Unless the user explicitly asks for the script, code, or implementation details, '
|
||||
'do not include the generated script in the final answer. '
|
||||
'A default workspace is mounted at /workspace for file tasks.'
|
||||
)
|
||||
),
|
||||
),
|
||||
skill_mgr=SimpleNamespace(
|
||||
get_skills_for_pipeline=AsyncMock(return_value=[]),
|
||||
detect_skill_activation=AsyncMock(return_value=None),
|
||||
build_activation_prompt=Mock(return_value=None),
|
||||
),
|
||||
)
|
||||
|
||||
runner = LocalAgentRunner(app, pipeline_config={})
|
||||
query = make_query()
|
||||
|
||||
results = [message async for message in runner.run(query)]
|
||||
|
||||
assert [message.role for message in results] == ['assistant', 'tool', 'assistant']
|
||||
assert results[-1].content == 'The average is 2.5.'
|
||||
|
||||
tool_manager.execute_func_call.assert_awaited_once()
|
||||
tool_name, tool_parameters = tool_manager.execute_func_call.await_args.args[:2]
|
||||
assert tool_name == 'exec'
|
||||
assert 'print(sum(nums) / len(nums))' in tool_parameters['command']
|
||||
|
||||
first_request = provider.requests[0]
|
||||
assert any(
|
||||
message.role == 'system'
|
||||
and 'exec' in str(message.content)
|
||||
and 'exact calculations' in str(message.content)
|
||||
and 'Unless the user explicitly asks for the script' in str(message.content)
|
||||
and '/workspace' in str(message.content)
|
||||
for message in first_request['messages']
|
||||
)
|
||||
assert [tool.name for tool in first_request['funcs']] == ['exec']
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_localagent_streaming_tool_error_yields_message_chunks():
|
||||
provider = RecordingStreamProvider()
|
||||
model = SimpleNamespace(
|
||||
provider=provider,
|
||||
model_entity=SimpleNamespace(
|
||||
uuid='test-model-uuid',
|
||||
name='test-model',
|
||||
abilities=['func_call'],
|
||||
extra_args={},
|
||||
),
|
||||
)
|
||||
|
||||
adapter = AsyncMock()
|
||||
adapter.is_stream_output_supported = AsyncMock(return_value=True)
|
||||
|
||||
query = make_query()
|
||||
query.adapter = adapter
|
||||
|
||||
app = SimpleNamespace(
|
||||
logger=Mock(),
|
||||
model_mgr=SimpleNamespace(get_model_by_uuid=AsyncMock(return_value=model)),
|
||||
tool_mgr=SimpleNamespace(execute_func_call=AsyncMock(side_effect=RuntimeError('boom'))),
|
||||
rag_mgr=SimpleNamespace(),
|
||||
box_service=SimpleNamespace(
|
||||
get_system_guidance=Mock(return_value='sandbox guidance'),
|
||||
),
|
||||
skill_mgr=SimpleNamespace(
|
||||
get_skills_for_pipeline=AsyncMock(return_value=[]),
|
||||
detect_skill_activation=AsyncMock(return_value=None),
|
||||
build_activation_prompt=Mock(return_value=None),
|
||||
),
|
||||
)
|
||||
|
||||
runner = LocalAgentRunner(app, pipeline_config={})
|
||||
|
||||
results = [message async for message in runner.run(query)]
|
||||
|
||||
assert all(isinstance(message, provider_message.MessageChunk) for message in results)
|
||||
assert any(message.role == 'tool' and message.content == 'err: boom' for message in results)
|
||||
@@ -0,0 +1,957 @@
|
||||
"""Tests for MCP Box integration: path rewriting, host_path inference, config model, payloads.
|
||||
|
||||
Uses importlib.util.spec_from_file_location to load mcp.py directly without
|
||||
triggering the circular import chain through the app module.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import importlib.util
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
from contextlib import asynccontextmanager
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Load mcp.py directly from file path, with stub dependencies
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _stub_module(fqn: str, attrs: dict | None = None, is_package: bool = False):
|
||||
"""Create or return a stub module and register it in sys.modules."""
|
||||
if fqn in sys.modules:
|
||||
mod = sys.modules[fqn]
|
||||
else:
|
||||
mod = types.ModuleType(fqn)
|
||||
mod.__spec__ = importlib.machinery.ModuleSpec(fqn, None, is_package=is_package)
|
||||
if is_package:
|
||||
mod.__path__ = []
|
||||
sys.modules[fqn] = mod
|
||||
parts = fqn.rsplit('.', 1)
|
||||
if len(parts) == 2 and parts[0] in sys.modules:
|
||||
setattr(sys.modules[parts[0]], parts[1], mod)
|
||||
if attrs:
|
||||
for k, v in attrs.items():
|
||||
setattr(mod, k, v)
|
||||
return mod
|
||||
|
||||
|
||||
@pytest.fixture(scope='module', autouse=True)
|
||||
def mcp_module():
|
||||
"""Load mcp.py with minimal stubs to avoid circular imports."""
|
||||
saved = {}
|
||||
|
||||
def _save_and_stub(name, attrs=None, is_package=False):
|
||||
saved[name] = sys.modules.get(name)
|
||||
# Don't overwrite modules that already exist (from other test modules)
|
||||
if name in sys.modules:
|
||||
return
|
||||
_stub_module(name, attrs, is_package)
|
||||
|
||||
# Stub entire dependency chains as packages / modules
|
||||
_save_and_stub('langbot_plugin', is_package=True)
|
||||
_save_and_stub('langbot_plugin.api', is_package=True)
|
||||
_save_and_stub('langbot_plugin.api.entities', is_package=True)
|
||||
_save_and_stub('langbot_plugin.api.entities.events', is_package=True)
|
||||
_save_and_stub('langbot_plugin.api.entities.events.pipeline_query', {})
|
||||
_save_and_stub('langbot_plugin.api.entities.builtin', is_package=True)
|
||||
_save_and_stub('langbot_plugin.api.entities.builtin.resource', is_package=True)
|
||||
_save_and_stub(
|
||||
'langbot_plugin.api.entities.builtin.resource.tool',
|
||||
{
|
||||
'LLMTool': type('LLMTool', (), {}),
|
||||
},
|
||||
)
|
||||
_save_and_stub('langbot_plugin.api.entities.builtin.provider', is_package=True)
|
||||
_save_and_stub('langbot_plugin.api.entities.builtin.provider.message', {})
|
||||
_save_and_stub('sqlalchemy', {'select': Mock()})
|
||||
_save_and_stub('httpx', {'AsyncClient': Mock()})
|
||||
_save_and_stub('mcp', {'ClientSession': Mock, 'StdioServerParameters': Mock}, is_package=True)
|
||||
_save_and_stub('mcp.client', is_package=True)
|
||||
_save_and_stub('mcp.client.stdio', {'stdio_client': Mock()})
|
||||
_save_and_stub('mcp.client.sse', {'sse_client': Mock()})
|
||||
_save_and_stub('mcp.client.streamable_http', {'streamable_http_client': Mock()})
|
||||
_save_and_stub('mcp.client.websocket', {'websocket_client': Mock()})
|
||||
|
||||
# Stub the provider.tools.loader (source of circular import)
|
||||
_save_and_stub('langbot', is_package=True)
|
||||
_save_and_stub('langbot.pkg', is_package=True)
|
||||
_save_and_stub('langbot.pkg.provider', is_package=True)
|
||||
_save_and_stub('langbot.pkg.provider.tools', is_package=True)
|
||||
_save_and_stub(
|
||||
'langbot.pkg.provider.tools.loader',
|
||||
{
|
||||
'ToolLoader': type('ToolLoader', (), {'__init__': lambda self, ap: None}),
|
||||
},
|
||||
)
|
||||
_save_and_stub('langbot.pkg.provider.tools.loaders', is_package=True)
|
||||
_save_and_stub('langbot.pkg.core', is_package=True)
|
||||
_save_and_stub('langbot.pkg.core.app', {'Application': type('Application', (), {})})
|
||||
_save_and_stub('langbot.pkg.entity', is_package=True)
|
||||
_save_and_stub('langbot.pkg.entity.persistence', is_package=True)
|
||||
_save_and_stub('langbot.pkg.entity.persistence.mcp', {})
|
||||
|
||||
# box models
|
||||
import enum as _enum
|
||||
|
||||
class _BPS(str, _enum.Enum):
|
||||
RUNNING = 'running'
|
||||
EXITED = 'exited'
|
||||
|
||||
_save_and_stub('langbot_plugin.box', is_package=True)
|
||||
_save_and_stub('langbot_plugin.box.models', {'BoxManagedProcessStatus': _BPS})
|
||||
|
||||
# Now load mcp.py via spec_from_file_location
|
||||
mod_fqn = 'langbot.pkg.provider.tools.loaders.mcp'
|
||||
sys.modules.pop(mod_fqn, None)
|
||||
mcp_path = os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
'..',
|
||||
'..',
|
||||
'..',
|
||||
'src',
|
||||
'langbot',
|
||||
'pkg',
|
||||
'provider',
|
||||
'tools',
|
||||
'loaders',
|
||||
'mcp.py',
|
||||
)
|
||||
mcp_path = os.path.normpath(mcp_path)
|
||||
pkg_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(mcp_path))))
|
||||
sys.modules['langbot.pkg'].__path__ = [pkg_root]
|
||||
sys.modules['langbot.pkg.provider.tools.loaders'].__path__ = [os.path.dirname(mcp_path)]
|
||||
spec = importlib.util.spec_from_file_location(mod_fqn, mcp_path)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
sys.modules[mod_fqn] = mod
|
||||
spec.loader.exec_module(mod)
|
||||
|
||||
yield mod
|
||||
|
||||
# Cleanup
|
||||
sys.modules.pop(mod_fqn, None)
|
||||
sys.modules.pop('langbot.pkg.provider.tools.loaders.mcp_stdio', None)
|
||||
sys.modules.pop('langbot.pkg.box.workspace', None)
|
||||
for name in reversed(list(saved)):
|
||||
if saved[name] is None:
|
||||
sys.modules.pop(name, None)
|
||||
else:
|
||||
sys.modules[name] = saved[name]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_ap():
|
||||
ap = Mock()
|
||||
ap.logger = Mock()
|
||||
ap.box_service = Mock()
|
||||
return ap
|
||||
|
||||
|
||||
def _make_session(mcp_module, server_config: dict, ap=None):
|
||||
if ap is None:
|
||||
ap = _make_ap()
|
||||
return mcp_module.RuntimeMCPSession(
|
||||
server_name=server_config.get('name', 'test-server'),
|
||||
server_config=server_config,
|
||||
enable=True,
|
||||
ap=ap,
|
||||
)
|
||||
|
||||
|
||||
# ── MCPServerBoxConfig ──────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestMCPServerBoxConfig:
|
||||
def test_default_values(self, mcp_module):
|
||||
cfg = mcp_module.MCPServerBoxConfig.model_validate({})
|
||||
assert cfg.image is None
|
||||
assert cfg.network == 'on'
|
||||
assert cfg.host_path is None
|
||||
assert cfg.host_path_mode == 'ro'
|
||||
assert cfg.env == {}
|
||||
assert cfg.startup_timeout_sec == 300
|
||||
assert cfg.cpus is None
|
||||
assert cfg.memory_mb is None
|
||||
assert cfg.pids_limit is None
|
||||
assert cfg.read_only_rootfs is None
|
||||
|
||||
def test_custom_values(self, mcp_module):
|
||||
cfg = mcp_module.MCPServerBoxConfig.model_validate(
|
||||
{
|
||||
'image': 'node:20',
|
||||
'network': 'on',
|
||||
'host_path': '/home/user/mcp',
|
||||
'host_path_mode': 'rw',
|
||||
'env': {'FOO': 'bar'},
|
||||
'startup_timeout_sec': 60,
|
||||
'cpus': 2.0,
|
||||
'memory_mb': 1024,
|
||||
'pids_limit': 256,
|
||||
'read_only_rootfs': False,
|
||||
}
|
||||
)
|
||||
assert cfg.image == 'node:20'
|
||||
assert cfg.network == 'on'
|
||||
assert cfg.cpus == 2.0
|
||||
assert cfg.memory_mb == 1024
|
||||
|
||||
def test_extra_fields_ignored(self, mcp_module):
|
||||
cfg = mcp_module.MCPServerBoxConfig.model_validate(
|
||||
{
|
||||
'image': 'node:20',
|
||||
'unknown_field': 'whatever',
|
||||
}
|
||||
)
|
||||
assert cfg.image == 'node:20'
|
||||
assert not hasattr(cfg, 'unknown_field')
|
||||
|
||||
|
||||
# ── Path Rewriting ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestRewritePath:
|
||||
def test_no_host_path_returns_unchanged(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
assert s._rewrite_path('/some/path', None) == '/some/path'
|
||||
|
||||
def test_empty_path_returns_empty(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
assert s._rewrite_path('', '/home/user/mcp') == ''
|
||||
|
||||
def test_prefix_match_rewrites(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
result = s._rewrite_path('/home/user/mcp/server.py', '/home/user/mcp')
|
||||
assert result == '/workspace/server.py'
|
||||
|
||||
def test_exact_match_rewrites_to_workspace(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
result = s._rewrite_path('/home/user/mcp', '/home/user/mcp')
|
||||
assert result == '/workspace'
|
||||
|
||||
def test_non_matching_path_unchanged(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
result = s._rewrite_path('/opt/other/server.py', '/home/user/mcp')
|
||||
assert result == '/opt/other/server.py'
|
||||
|
||||
def test_similar_prefix_not_rewritten(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
result = s._rewrite_path('/home/user/mcp-other/file.py', '/home/user/mcp')
|
||||
assert result == '/home/user/mcp-other/file.py'
|
||||
|
||||
def test_nested_subpath_rewrites(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
result = s._rewrite_path('/home/user/mcp/src/lib/main.py', '/home/user/mcp')
|
||||
assert result == '/workspace/src/lib/main.py'
|
||||
|
||||
|
||||
# ── host_path Inference ─────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestInferHostPath:
|
||||
def test_no_absolute_paths_returns_none(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': ['server.py'],
|
||||
},
|
||||
)
|
||||
assert s._infer_host_path() is None
|
||||
|
||||
def test_nonexistent_path_returns_none(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': '/nonexistent/path/to/python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
assert s._infer_host_path() is None
|
||||
|
||||
def test_existing_absolute_path_infers_directory(self, mcp_module):
|
||||
with tempfile.NamedTemporaryFile(suffix='.py') as f:
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [f.name],
|
||||
},
|
||||
)
|
||||
result = s._infer_host_path()
|
||||
assert result is not None
|
||||
assert result == os.path.dirname(os.path.realpath(f.name))
|
||||
|
||||
|
||||
# ── Build Box Session Payload ───────────────────────────────────────
|
||||
|
||||
|
||||
class TestBuildBoxSessionPayload:
|
||||
def test_minimal_config(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
payload = s._build_box_session_payload('session-123')
|
||||
assert payload['session_id'] == 'session-123'
|
||||
assert payload['workdir'] == '/workspace'
|
||||
assert payload['env'] == {}
|
||||
assert 'host_path' not in payload
|
||||
|
||||
def test_with_host_path(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
'box': {'host_path': '/home/user/mcp', 'host_path_mode': 'ro'},
|
||||
},
|
||||
)
|
||||
payload = s._build_box_session_payload('session-123')
|
||||
assert payload['host_path'] == '/home/user/mcp'
|
||||
assert payload['host_path_mode'] == 'ro'
|
||||
|
||||
def test_optional_fields_included_when_set(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
'box': {'image': 'node:20', 'cpus': 2.0, 'memory_mb': 1024, 'pids_limit': 256},
|
||||
},
|
||||
)
|
||||
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['pids_limit'] == 256
|
||||
|
||||
def test_none_fields_excluded(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
payload = s._build_box_session_payload('session-123')
|
||||
assert 'image' not in payload
|
||||
assert 'cpus' not in payload
|
||||
|
||||
|
||||
# ── Build Box Process Payload ───────────────────────────────────────
|
||||
|
||||
|
||||
class TestBuildBoxProcessPayload:
|
||||
def test_basic_payload(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': ['server.py'],
|
||||
'env': {'KEY': 'val'},
|
||||
},
|
||||
)
|
||||
payload = s._build_box_process_payload()
|
||||
assert payload['command'] == 'python'
|
||||
assert payload['args'] == ['server.py']
|
||||
assert payload['env'] == {'KEY': 'val'}
|
||||
assert payload['cwd'] == '/workspace'
|
||||
|
||||
def test_path_rewriting_applied(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': '/home/user/mcp/venv/bin/python',
|
||||
'args': ['/home/user/mcp/server.py', '--config', '/home/user/mcp/config.json'],
|
||||
'env': {},
|
||||
'box': {'host_path': '/home/user/mcp'},
|
||||
},
|
||||
)
|
||||
payload = s._build_box_process_payload()
|
||||
# venv python is replaced with plain 'python' (deps installed in-container)
|
||||
assert payload['command'] == 'python'
|
||||
assert payload['args'] == ['/workspace/server.py', '--config', '/workspace/config.json']
|
||||
|
||||
def test_non_matching_args_not_rewritten(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': ['/opt/other/server.py', '--flag'],
|
||||
'env': {},
|
||||
'box': {'host_path': '/home/user/mcp'},
|
||||
},
|
||||
)
|
||||
payload = s._build_box_process_payload()
|
||||
assert payload['command'] == 'python'
|
||||
assert payload['args'] == ['/opt/other/server.py', '--flag']
|
||||
|
||||
|
||||
# ── Python Workspace Preparation ────────────────────────────────────
|
||||
|
||||
|
||||
class TestPythonWorkspacePreparation:
|
||||
def test_requirements_workspace_uses_venv_bootstrap(self, mcp_module, tmp_path):
|
||||
host_path = tmp_path / 'mcp-source'
|
||||
host_path.mkdir()
|
||||
(host_path / 'requirements.txt').write_text('mcp==1.26.0\n', encoding='utf-8')
|
||||
|
||||
command = mcp_module.BoxStdioSessionRuntime.detect_install_command(
|
||||
str(host_path),
|
||||
'/workspace/.mcp/u1/workspace',
|
||||
)
|
||||
|
||||
assert command is not None
|
||||
assert '_LB_SYSTEM_PYTHON="$(command -v python3 || command -v python || true)"' in command
|
||||
assert '"$_LB_SYSTEM_PYTHON" -m venv "$_LB_VENV_DIR"' in command
|
||||
assert 'python -m pip install -r "/workspace/.mcp/u1/workspace/requirements.txt"' in command
|
||||
assert 'pip install --no-cache-dir -r' not in command
|
||||
|
||||
def test_staging_refresh_removes_stale_source_files_but_preserves_runtime_dirs(self, mcp_module, tmp_path):
|
||||
source = tmp_path / 'source'
|
||||
source.mkdir()
|
||||
(source / 'server.py').write_text('print("new")\n', encoding='utf-8')
|
||||
(source / 'requirements.txt').write_text('mcp==1.26.0\n', encoding='utf-8')
|
||||
(source / '.env').write_text('TOKEN=new\n', encoding='utf-8')
|
||||
|
||||
process_root = tmp_path / 'shared' / '.mcp' / 'u1'
|
||||
workspace = process_root / 'workspace'
|
||||
(workspace / '.venv' / 'bin').mkdir(parents=True)
|
||||
(workspace / '.venv' / 'bin' / 'python').write_text('', encoding='utf-8')
|
||||
(workspace / '.langbot').mkdir()
|
||||
(workspace / '.langbot' / 'python-env.lock').mkdir()
|
||||
(workspace / '.env').write_text('TOKEN=old\n', encoding='utf-8')
|
||||
(workspace / 'server.py').write_text('print("old")\n', encoding='utf-8')
|
||||
(workspace / 'removed.py').write_text('stale\n', encoding='utf-8')
|
||||
(workspace / 'removed_dir').mkdir()
|
||||
(workspace / 'removed_dir' / 'old.txt').write_text('stale\n', encoding='utf-8')
|
||||
|
||||
mcp_module.BoxStdioSessionRuntime._copy_workspace_tree(str(source), str(process_root), str(workspace))
|
||||
|
||||
assert (workspace / 'server.py').read_text(encoding='utf-8') == 'print("new")\n'
|
||||
assert (workspace / 'requirements.txt').read_text(encoding='utf-8') == 'mcp==1.26.0\n'
|
||||
assert (workspace / '.env').read_text(encoding='utf-8') == 'TOKEN=new\n'
|
||||
assert not (workspace / 'removed.py').exists()
|
||||
assert not (workspace / 'removed_dir').exists()
|
||||
assert (workspace / '.venv' / 'bin' / 'python').exists()
|
||||
assert (workspace / '.langbot' / 'python-env.lock').is_dir()
|
||||
|
||||
def test_staging_refresh_ignores_unlink_race(self, mcp_module, tmp_path, monkeypatch):
|
||||
mcp_stdio_module = sys.modules['langbot.pkg.provider.tools.loaders.mcp_stdio']
|
||||
|
||||
source = tmp_path / 'source'
|
||||
source.mkdir()
|
||||
(source / 'server.py').write_text('print("new")\n', encoding='utf-8')
|
||||
|
||||
process_root = tmp_path / 'shared' / '.mcp' / 'u1'
|
||||
workspace = process_root / 'workspace'
|
||||
workspace.mkdir(parents=True)
|
||||
stale_file = workspace / 'removed.py'
|
||||
stale_file.write_text('stale\n', encoding='utf-8')
|
||||
|
||||
real_unlink = os.unlink
|
||||
|
||||
def unlink_with_race(path):
|
||||
if os.fspath(path) == str(stale_file):
|
||||
real_unlink(path)
|
||||
raise FileNotFoundError(path)
|
||||
real_unlink(path)
|
||||
|
||||
monkeypatch.setattr(mcp_stdio_module.os, 'unlink', unlink_with_race)
|
||||
|
||||
mcp_module.BoxStdioSessionRuntime._copy_workspace_tree(str(source), str(process_root), str(workspace))
|
||||
|
||||
assert not stale_file.exists()
|
||||
assert (workspace / 'server.py').read_text(encoding='utf-8') == 'print("new")\n'
|
||||
|
||||
|
||||
# ── get_runtime_info_dict ───────────────────────────────────────────
|
||||
|
||||
|
||||
class TestGetRuntimeInfoDict:
|
||||
def test_non_stdio_session(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'test-uuid',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
info = s.get_runtime_info_dict()
|
||||
assert info['status'] == 'connecting'
|
||||
assert 'box_session_id' not in info
|
||||
|
||||
def test_runtime_tools_include_parameters(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'test-uuid',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
s.functions = [
|
||||
SimpleNamespace(
|
||||
name='create-service',
|
||||
description='Create a service',
|
||||
parameters={
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'project_id': {'type': 'string'},
|
||||
},
|
||||
'required': ['project_id'],
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
info = s.get_runtime_info_dict()
|
||||
|
||||
assert info['tools'][0]['parameters']['properties']['project_id']['type'] == 'string'
|
||||
assert info['tools'][0]['parameters']['required'] == ['project_id']
|
||||
|
||||
def test_stdio_session_includes_box_info(self, mcp_module):
|
||||
ap = _make_ap()
|
||||
ap.box_service.available = True
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'test-uuid',
|
||||
'mode': 'stdio',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
ap=ap,
|
||||
)
|
||||
info = s.get_runtime_info_dict()
|
||||
assert info['box_session_id'] == 'mcp-shared'
|
||||
assert info['box_enabled'] is True
|
||||
|
||||
def test_transient_test_shares_session_but_isolated_by_process(self, mcp_module):
|
||||
"""A transient config-page "test" now shares the same 'mcp-shared' Box
|
||||
session as live servers (so a test reuses the running container / live
|
||||
process instead of a cold per-test session bootstrap). Isolation is at
|
||||
the PROCESS level: the test runs under its own process_id and only ever
|
||||
stops that process_id, so it cannot disturb another server's live
|
||||
process or the shared session itself."""
|
||||
ap = _make_ap()
|
||||
ap.box_service.available = True
|
||||
transient = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'gen-uuid-123',
|
||||
'mode': 'stdio',
|
||||
'command': 'uvx',
|
||||
'args': ['mcp-server-time'],
|
||||
'_transient': True,
|
||||
},
|
||||
ap=ap,
|
||||
)
|
||||
live = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'time',
|
||||
'uuid': 'real-uuid',
|
||||
'mode': 'stdio',
|
||||
'command': 'uvx',
|
||||
'args': ['mcp-server-time'],
|
||||
},
|
||||
ap=ap,
|
||||
)
|
||||
assert transient.is_transient is True
|
||||
assert live.is_transient is False
|
||||
# Both share ONE Box session ...
|
||||
assert transient._build_box_session_id() == 'mcp-shared'
|
||||
assert live._build_box_session_id() == 'mcp-shared'
|
||||
assert transient._build_box_session_id() == live._build_box_session_id()
|
||||
# ... but are isolated by distinct process_ids within that session.
|
||||
assert transient._box_stdio_runtime.process_id != live._box_stdio_runtime.process_id
|
||||
|
||||
def test_stdio_session_refuses_when_box_unavailable(self, mcp_module):
|
||||
"""Policy: when Box is configured but unavailable (disabled in config
|
||||
OR connection failed), stdio MCP servers are NOT treated as box-stdio.
|
||||
``_init_stdio_python_server`` will raise a clear refusal at start
|
||||
time; until then, the runtime info simply omits box_session_id so the
|
||||
UI can render the disabled state cleanly."""
|
||||
ap = _make_ap()
|
||||
ap.box_service.available = False
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'test-uuid',
|
||||
'mode': 'stdio',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
ap=ap,
|
||||
)
|
||||
info = s.get_runtime_info_dict()
|
||||
assert 'box_session_id' not in info
|
||||
assert 'box_enabled' not in info
|
||||
|
||||
def test_stdio_session_without_box_service_uses_local_stdio(self, mcp_module):
|
||||
ap = _make_ap()
|
||||
del ap.box_service
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'test-uuid',
|
||||
'mode': 'stdio',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
ap=ap,
|
||||
)
|
||||
info = s.get_runtime_info_dict()
|
||||
assert 'box_session_id' not in info
|
||||
|
||||
|
||||
# ── Box config parsing ──────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestBoxConfigParsing:
|
||||
def test_box_config_parsed_from_server_config(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
'box': {'image': 'node:20', 'host_path': '/home/user/mcp'},
|
||||
},
|
||||
)
|
||||
assert isinstance(s.box_config, mcp_module.MCPServerBoxConfig)
|
||||
assert s.box_config.image == 'node:20'
|
||||
assert s.box_config.host_path == '/home/user/mcp'
|
||||
|
||||
def test_missing_box_key_uses_defaults(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
assert isinstance(s.box_config, mcp_module.MCPServerBoxConfig)
|
||||
assert s.box_config.image is None
|
||||
assert s.box_config.host_path_mode == 'ro'
|
||||
|
||||
|
||||
@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']
|
||||
|
||||
class FakeClientSession:
|
||||
def __init__(self, *_args):
|
||||
pass
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
async def initialize(self):
|
||||
return None
|
||||
|
||||
@asynccontextmanager
|
||||
async def fake_websocket_client(_url: str):
|
||||
yield ('read-stream', 'write-stream')
|
||||
|
||||
mcp_stdio_module.ClientSession = FakeClientSession
|
||||
mcp_stdio_module.websocket_client = fake_websocket_client
|
||||
|
||||
ap = _make_ap()
|
||||
ap.box_service.available = True
|
||||
ap.box_service.default_workspace = str(tmp_path / 'shared-box-workspace')
|
||||
ap.box_service.create_session = AsyncMock(return_value={})
|
||||
ap.box_service.build_spec = Mock(return_value='validated-spec')
|
||||
ap.box_service.client = SimpleNamespace(
|
||||
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')
|
||||
|
||||
host_path = tmp_path / 'mcp-source'
|
||||
host_path.mkdir()
|
||||
server_file = host_path / 'server.py'
|
||||
server_file.write_text('print("hello")\n', encoding='utf-8')
|
||||
|
||||
session = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'stdio',
|
||||
'command': str(host_path / '.venv' / 'bin' / 'python'),
|
||||
'args': [str(server_file)],
|
||||
'box': {'host_path': str(host_path)},
|
||||
},
|
||||
ap=ap,
|
||||
)
|
||||
|
||||
await session._init_box_stdio_server()
|
||||
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 session_payload['session_id'] == 'mcp-shared'
|
||||
assert 'host_path' not in session_payload
|
||||
assert ap.box_service.build_spec.call_count == 1
|
||||
assert ap.box_service.build_spec.call_args.kwargs.get('skip_host_mount_validation', False) is False
|
||||
assert ap.box_service.build_spec.call_args.args[0]['host_path'] == str(host_path)
|
||||
|
||||
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 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'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stdio_handshake_raises_coldstart_retry_while_process_alive(mcp_module, tmp_path, monkeypatch):
|
||||
"""During a slow (npx) cold start the handshake fails while the managed
|
||||
process is still alive. initialize() must raise _ColdStartRetry (so the
|
||||
outer lifecycle loop reuses the live process and retries without stopping it
|
||||
or consuming the fatal budget), NOT a fatal error."""
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
mcp_stdio_module = sys.modules['langbot.pkg.provider.tools.loaders.mcp_stdio']
|
||||
|
||||
class ColdClientSession:
|
||||
def __init__(self, *_args):
|
||||
pass
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
async def initialize(self):
|
||||
# Process still cold-starting: handshake fails.
|
||||
raise Exception('Connection closed')
|
||||
|
||||
@asynccontextmanager
|
||||
async def fake_websocket_client(_url: str):
|
||||
yield ('read-stream', 'write-stream')
|
||||
|
||||
monkeypatch.setattr(mcp_stdio_module, 'ClientSession', ColdClientSession)
|
||||
monkeypatch.setattr(mcp_stdio_module, 'websocket_client', fake_websocket_client)
|
||||
monkeypatch.setattr(mcp_stdio_module, '_HANDSHAKE_ATTEMPT_TIMEOUT_SEC', 1.0, raising=False)
|
||||
|
||||
ap = _make_ap()
|
||||
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')
|
||||
|
||||
session = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'slow',
|
||||
'uuid': 'slow-uuid',
|
||||
'mode': 'stdio',
|
||||
'command': 'npx',
|
||||
'args': ['-y', 'some-mcp'],
|
||||
},
|
||||
ap=ap,
|
||||
)
|
||||
|
||||
# Process is NOT exited (still cold-starting) and not yet running for reuse.
|
||||
async def _not_exited():
|
||||
return False
|
||||
|
||||
session._box_stdio_runtime._managed_process_has_exited = _not_exited
|
||||
|
||||
async def _not_running():
|
||||
return False
|
||||
|
||||
session._box_stdio_runtime._managed_process_is_running = _not_running
|
||||
|
||||
with pytest.raises(mcp_stdio_module._ColdStartRetry):
|
||||
await session._init_box_stdio_server()
|
||||
|
||||
# Process was started exactly once (the retry will reuse it, not rebuild).
|
||||
assert ap.box_service.start_managed_process.await_count == 1
|
||||
await session.exit_stack.aclose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stdio_handshake_raises_fatal_when_process_exited(mcp_module, tmp_path, monkeypatch):
|
||||
"""If the handshake fails AND the process has definitively exited, that is a
|
||||
real failure — initialize() must NOT swallow it as a cold-start retry."""
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
mcp_stdio_module = sys.modules['langbot.pkg.provider.tools.loaders.mcp_stdio']
|
||||
|
||||
class DeadClientSession:
|
||||
def __init__(self, *_args):
|
||||
pass
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
async def initialize(self):
|
||||
raise Exception('Connection closed')
|
||||
|
||||
@asynccontextmanager
|
||||
async def fake_websocket_client(_url: str):
|
||||
yield ('read-stream', 'write-stream')
|
||||
|
||||
monkeypatch.setattr(mcp_stdio_module, 'ClientSession', DeadClientSession)
|
||||
monkeypatch.setattr(mcp_stdio_module, 'websocket_client', fake_websocket_client)
|
||||
monkeypatch.setattr(mcp_stdio_module, '_HANDSHAKE_ATTEMPT_TIMEOUT_SEC', 1.0, raising=False)
|
||||
|
||||
ap = _make_ap()
|
||||
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')
|
||||
|
||||
session = _make_session(
|
||||
mcp_module,
|
||||
{'name': 'dead', 'uuid': 'dead-uuid', 'mode': 'stdio', 'command': 'npx', 'args': ['-y', 'x']},
|
||||
ap=ap,
|
||||
)
|
||||
|
||||
async def _exited():
|
||||
return True
|
||||
|
||||
session._box_stdio_runtime._managed_process_has_exited = _exited
|
||||
|
||||
async def _not_running():
|
||||
return False
|
||||
|
||||
session._box_stdio_runtime._managed_process_is_running = _not_running
|
||||
|
||||
with pytest.raises(Exception) as ei:
|
||||
await session._init_box_stdio_server()
|
||||
assert not isinstance(ei.value, mcp_stdio_module._ColdStartRetry)
|
||||
await session.exit_stack.aclose()
|
||||
@@ -0,0 +1,244 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from contextlib import asynccontextmanager
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, cast
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from aiohttp import web
|
||||
from mcp import types as mcp_types
|
||||
|
||||
from langbot.pkg.provider.tools.loaders.mcp import RuntimeMCPSession
|
||||
|
||||
|
||||
class _TransportProbe:
|
||||
def __init__(self, streamable_status: int | None) -> None:
|
||||
self.streamable_status = streamable_status
|
||||
self.streamable_posts = 0
|
||||
self.streamable_messages: list[str] = []
|
||||
self.sse_gets = 0
|
||||
self.sse_messages: list[str] = []
|
||||
self.streamable_request_started = asyncio.Event()
|
||||
self.release_streamable_request = asyncio.Event()
|
||||
self._sse_response: web.StreamResponse | None = None
|
||||
|
||||
async def handle_mcp_endpoint(self, request: web.Request) -> web.StreamResponse:
|
||||
if request.method == 'POST':
|
||||
self.streamable_posts += 1
|
||||
self.streamable_request_started.set()
|
||||
if self.streamable_status is None:
|
||||
await self.release_streamable_request.wait()
|
||||
return web.Response(status=204)
|
||||
if self.streamable_status == 200:
|
||||
message = await request.json()
|
||||
method = message.get('method', '')
|
||||
self.streamable_messages.append(method)
|
||||
if method == 'initialize':
|
||||
return web.json_response(
|
||||
{
|
||||
'jsonrpc': '2.0',
|
||||
'id': message['id'],
|
||||
'result': {
|
||||
'protocolVersion': mcp_types.LATEST_PROTOCOL_VERSION,
|
||||
'capabilities': {'tools': {}},
|
||||
'serverInfo': {'name': 'streamable-test', 'version': '1.0.0'},
|
||||
},
|
||||
}
|
||||
)
|
||||
if method == 'tools/list':
|
||||
return web.json_response(
|
||||
{
|
||||
'jsonrpc': '2.0',
|
||||
'id': message['id'],
|
||||
'result': {
|
||||
'tools': [
|
||||
{
|
||||
'name': 'echo',
|
||||
'description': 'Echo test input',
|
||||
'inputSchema': {'type': 'object'},
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
)
|
||||
return web.Response(status=202)
|
||||
return web.Response(status=self.streamable_status)
|
||||
|
||||
self.sse_gets += 1
|
||||
response = web.StreamResponse(
|
||||
status=200,
|
||||
headers={
|
||||
'Content-Type': 'text/event-stream',
|
||||
'Cache-Control': 'no-cache',
|
||||
},
|
||||
)
|
||||
await response.prepare(request)
|
||||
self._sse_response = response
|
||||
await response.write(b'event: endpoint\ndata: /messages?session_id=test-session\n\n')
|
||||
try:
|
||||
while request.transport is not None and not request.transport.is_closing():
|
||||
await asyncio.sleep(0.05)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
return response
|
||||
|
||||
async def handle_sse_message(self, request: web.Request) -> web.Response:
|
||||
message = await request.json()
|
||||
method = message.get('method', '')
|
||||
self.sse_messages.append(method)
|
||||
|
||||
if method == 'initialize':
|
||||
response_message = {
|
||||
'jsonrpc': '2.0',
|
||||
'id': message['id'],
|
||||
'result': {
|
||||
'protocolVersion': mcp_types.LATEST_PROTOCOL_VERSION,
|
||||
'capabilities': {},
|
||||
'serverInfo': {'name': 'legacy-sse-test', 'version': '1.0.0'},
|
||||
},
|
||||
}
|
||||
assert self._sse_response is not None
|
||||
payload = json.dumps(response_message, separators=(',', ':'))
|
||||
await self._sse_response.write(f'event: message\ndata: {payload}\n\n'.encode())
|
||||
|
||||
return web.Response(status=202)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _transport_server(streamable_status: int | None):
|
||||
probe = _TransportProbe(streamable_status)
|
||||
application = web.Application()
|
||||
application.router.add_route('*', '/mcp', probe.handle_mcp_endpoint)
|
||||
application.router.add_post('/messages', probe.handle_sse_message)
|
||||
runner = web.AppRunner(application, shutdown_timeout=0.1)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, '127.0.0.1', 0)
|
||||
await site.start()
|
||||
server = cast(asyncio.Server, site._server)
|
||||
port = server.sockets[0].getsockname()[1]
|
||||
try:
|
||||
yield probe, f'http://127.0.0.1:{port}/mcp'
|
||||
finally:
|
||||
await runner.cleanup()
|
||||
|
||||
|
||||
def _session(url: str, *, timeout: float = 2) -> RuntimeMCPSession:
|
||||
app = cast(Any, SimpleNamespace(logger=Mock()))
|
||||
return RuntimeMCPSession(
|
||||
'remote-transport-test',
|
||||
{'uuid': 'srv-1', 'mode': 'remote', 'url': url, 'timeout': timeout},
|
||||
True,
|
||||
app,
|
||||
)
|
||||
|
||||
|
||||
def _contains_http_status(exc: BaseException, status_code: int) -> bool:
|
||||
return any(
|
||||
isinstance(leaf, httpx.HTTPStatusError) and leaf.response.status_code == status_code
|
||||
for leaf in RuntimeMCPSession._iter_exception_leaves(exc)
|
||||
)
|
||||
|
||||
|
||||
async def _close_session(session: RuntimeMCPSession) -> None:
|
||||
await session.exit_stack.aclose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_remote_transport_real_streamable_http_success_keeps_session_usable():
|
||||
async with _transport_server(200) as (probe, url):
|
||||
session = _session(url)
|
||||
try:
|
||||
await session._init_remote_server()
|
||||
assert session.session is not None
|
||||
tools = await session.session.list_tools()
|
||||
assert [tool.name for tool in tools.tools] == ['echo']
|
||||
assert probe.streamable_posts >= 2
|
||||
assert probe.streamable_messages[:2] == ['initialize', 'notifications/initialized']
|
||||
assert 'tools/list' in probe.streamable_messages
|
||||
assert probe.sse_gets == 0
|
||||
finally:
|
||||
await _close_session(session)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('status_code', [400, 404, 405])
|
||||
async def test_remote_transport_real_streamable_http_error_falls_back_to_legacy_sse(status_code: int):
|
||||
async with _transport_server(status_code) as (probe, url):
|
||||
session = _session(url)
|
||||
try:
|
||||
await session._init_remote_server()
|
||||
assert session.session is not None
|
||||
assert probe.streamable_posts == 1
|
||||
assert probe.sse_gets == 1
|
||||
assert 'initialize' in probe.sse_messages
|
||||
finally:
|
||||
await _close_session(session)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('status_code', [401, 403, 406, 415, 429, 500])
|
||||
async def test_remote_transport_real_non_compatibility_error_does_not_fallback(status_code: int):
|
||||
async with _transport_server(status_code) as (probe, url):
|
||||
session = _session(url)
|
||||
try:
|
||||
with pytest.raises(BaseException) as exc_info:
|
||||
await session._init_remote_server()
|
||||
assert _contains_http_status(exc_info.value, status_code)
|
||||
assert probe.streamable_posts == 1
|
||||
assert probe.sse_gets == 0
|
||||
finally:
|
||||
await _close_session(session)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_remote_transport_real_timeout_does_not_fallback():
|
||||
async with _transport_server(None) as (probe, url):
|
||||
session = _session(url, timeout=0.05)
|
||||
try:
|
||||
with pytest.raises(BaseException) as exc_info:
|
||||
await session._init_remote_server()
|
||||
assert any(
|
||||
isinstance(leaf, httpx.TimeoutException)
|
||||
for leaf in RuntimeMCPSession._iter_exception_leaves(exc_info.value)
|
||||
)
|
||||
assert probe.streamable_posts == 1
|
||||
assert probe.sse_gets == 0
|
||||
finally:
|
||||
probe.release_streamable_request.set()
|
||||
await _close_session(session)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('error_type', [httpx.ConnectError, httpx.ConnectTimeout])
|
||||
async def test_remote_transport_connection_errors_do_not_fallback(error_type: type[httpx.RequestError]):
|
||||
request = httpx.Request('POST', 'https://unreachable.invalid/mcp')
|
||||
error = error_type('connection failed', request=request)
|
||||
session = _session(str(request.url))
|
||||
session._init_streamable_http_server = AsyncMock(side_effect=error)
|
||||
session._init_sse_server = AsyncMock()
|
||||
|
||||
with pytest.raises(type(error)) as exc_info:
|
||||
await session._init_remote_server()
|
||||
|
||||
assert exc_info.value is error
|
||||
session._init_sse_server.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_remote_transport_external_cancellation_is_not_converted_to_sse_fallback():
|
||||
async with _transport_server(None) as (probe, url):
|
||||
session = _session(url)
|
||||
task = asyncio.create_task(session._init_remote_server())
|
||||
await asyncio.wait_for(probe.streamable_request_started.wait(), timeout=2)
|
||||
task.cancel()
|
||||
try:
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert probe.sse_gets == 0
|
||||
finally:
|
||||
probe.release_streamable_request.set()
|
||||
await _close_session(session)
|
||||
@@ -0,0 +1,323 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from mcp import types as mcp_types
|
||||
|
||||
from langbot.pkg.provider.tools.loaders.mcp import (
|
||||
MCP_RESOURCE_CONTEXT_QUERY_KEY,
|
||||
MCP_RESOURCE_TRACE_QUERY_KEY,
|
||||
MCP_TOOL_LIST_RESOURCES,
|
||||
MCP_TOOL_READ_RESOURCE,
|
||||
MCPLoader,
|
||||
MCPSessionStatus,
|
||||
RuntimeMCPSession,
|
||||
)
|
||||
from langbot.pkg.telemetry import features as telemetry_features
|
||||
|
||||
|
||||
def _app() -> SimpleNamespace:
|
||||
return SimpleNamespace(logger=Mock())
|
||||
|
||||
|
||||
def _connected_session(
|
||||
*,
|
||||
name: str = 'docs',
|
||||
uuid: str = 'srv-1',
|
||||
resources: list[dict] | None = None,
|
||||
templates: list[dict] | None = None,
|
||||
) -> RuntimeMCPSession:
|
||||
session = RuntimeMCPSession(name, {'uuid': uuid, 'mode': 'remote'}, True, _app())
|
||||
session.status = MCPSessionStatus.CONNECTED
|
||||
session.session = SimpleNamespace(read_resource=AsyncMock())
|
||||
session.resources = resources or [
|
||||
{
|
||||
'uri': 'file:///README.md',
|
||||
'name': 'README.md',
|
||||
'title': '',
|
||||
'description': '',
|
||||
'mime_type': 'text/markdown',
|
||||
'size': None,
|
||||
'icons': [],
|
||||
'annotations': {},
|
||||
'_meta': {},
|
||||
}
|
||||
]
|
||||
session.resource_templates = templates or []
|
||||
return session
|
||||
|
||||
|
||||
def _query() -> SimpleNamespace:
|
||||
return SimpleNamespace(variables={})
|
||||
|
||||
|
||||
def _http_status_error(status_code: int) -> httpx.HTTPStatusError:
|
||||
request = httpx.Request('POST', 'https://example.com/mcp')
|
||||
response = httpx.Response(status_code, request=request)
|
||||
return httpx.HTTPStatusError(f'HTTP {status_code}', request=request, response=response)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_remote_transport_falls_back_to_sse_for_compatible_http_status_in_exception_group():
|
||||
session = RuntimeMCPSession(
|
||||
'remote',
|
||||
{'uuid': 'srv-1', 'mode': 'remote', 'url': 'https://example.com/mcp'},
|
||||
True,
|
||||
_app(),
|
||||
)
|
||||
session._init_streamable_http_server = AsyncMock(
|
||||
side_effect=ExceptionGroup('transport failed', [_http_status_error(405)])
|
||||
)
|
||||
session._init_sse_server = AsyncMock()
|
||||
|
||||
await session._init_remote_server()
|
||||
|
||||
session._init_streamable_http_server.assert_awaited_once()
|
||||
session._init_sse_server.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_remote_transport_does_not_fallback_for_auth_http_status():
|
||||
session = RuntimeMCPSession(
|
||||
'remote',
|
||||
{'uuid': 'srv-1', 'mode': 'remote', 'url': 'https://example.com/mcp'},
|
||||
True,
|
||||
_app(),
|
||||
)
|
||||
error = _http_status_error(403)
|
||||
session._init_streamable_http_server = AsyncMock(side_effect=error)
|
||||
session._init_sse_server = AsyncMock()
|
||||
|
||||
with pytest.raises(httpx.HTTPStatusError):
|
||||
await session._init_remote_server()
|
||||
|
||||
session._init_streamable_http_server.assert_awaited_once()
|
||||
session._init_sse_server.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_resource_envelope_truncates_caches_and_records_trace():
|
||||
session = _connected_session()
|
||||
session.session.read_resource.return_value = mcp_types.ReadResourceResult(
|
||||
contents=[
|
||||
mcp_types.TextResourceContents(
|
||||
uri='file:///README.md',
|
||||
mimeType='text/markdown',
|
||||
text='abcdef',
|
||||
)
|
||||
]
|
||||
)
|
||||
query = _query()
|
||||
|
||||
first = await session.read_resource_envelope(
|
||||
'file:///README.md',
|
||||
max_bytes=4,
|
||||
source='ui_preview',
|
||||
query=query,
|
||||
)
|
||||
second = await session.read_resource_envelope(
|
||||
'file:///README.md',
|
||||
max_bytes=4,
|
||||
source='agent_tool',
|
||||
query=query,
|
||||
)
|
||||
|
||||
assert first['contents'][0]['text'] == 'abcd'
|
||||
assert first['contents'][0]['bytes'] == 6
|
||||
assert first['truncated'] is True
|
||||
assert first['cache_hit'] is False
|
||||
assert second['cache_hit'] is True
|
||||
assert second['source'] == 'agent_tool'
|
||||
assert session.session.read_resource.await_count == 1
|
||||
|
||||
traces = query.variables[MCP_RESOURCE_TRACE_QUERY_KEY]
|
||||
assert [trace['source'] for trace in traces] == ['ui_preview', 'agent_tool']
|
||||
assert traces[1]['cache_hit'] is True
|
||||
assert query.variables[telemetry_features.FEATURES_KEY]['mcp_resource_reads'] == {
|
||||
'ui_preview': 1,
|
||||
'agent_tool': 1,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_resource_envelope_shares_byte_budget_across_text_contents():
|
||||
session = _connected_session()
|
||||
session.session.read_resource.return_value = mcp_types.ReadResourceResult(
|
||||
contents=[
|
||||
mcp_types.TextResourceContents(
|
||||
uri='file:///README.md#first',
|
||||
mimeType='text/plain',
|
||||
text='abc',
|
||||
),
|
||||
mcp_types.TextResourceContents(
|
||||
uri='file:///README.md#second',
|
||||
mimeType='text/plain',
|
||||
text='def',
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
envelope = await session.read_resource_envelope('file:///README.md', max_bytes=4)
|
||||
|
||||
assert [item['text'] for item in envelope['contents']] == ['abc', 'd']
|
||||
assert envelope['contents'][0]['truncated'] is False
|
||||
assert envelope['contents'][1]['truncated'] is True
|
||||
assert envelope['bytes'] == 6
|
||||
assert envelope['truncated'] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_resource_envelope_omits_binary_by_default():
|
||||
session = _connected_session(
|
||||
resources=[
|
||||
{
|
||||
'uri': 'file:///image.png',
|
||||
'name': 'image.png',
|
||||
'title': '',
|
||||
'description': '',
|
||||
'mime_type': 'image/png',
|
||||
'size': 4,
|
||||
'icons': [],
|
||||
'annotations': {},
|
||||
'_meta': {},
|
||||
}
|
||||
]
|
||||
)
|
||||
session.session.read_resource.return_value = mcp_types.ReadResourceResult(
|
||||
contents=[
|
||||
mcp_types.BlobResourceContents(
|
||||
uri='file:///image.png',
|
||||
mimeType='image/png',
|
||||
blob=base64.b64encode(b'\x00\x01\x02\x03').decode(),
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
envelope = await session.read_resource_envelope('file:///image.png')
|
||||
|
||||
content = envelope['contents'][0]
|
||||
assert content['type'] == 'blob'
|
||||
assert content['blob'] is None
|
||||
assert content['bytes'] == 4
|
||||
assert content['binary_omitted'] is True
|
||||
assert envelope['truncated'] is True
|
||||
assert envelope['warnings'] == ['Binary resource content omitted from response.']
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_resource_envelope_rejects_unlisted_uri():
|
||||
session = _connected_session()
|
||||
|
||||
with pytest.raises(ValueError, match='Resource URI is not available'):
|
||||
await session.read_resource_envelope('file:///secret.txt')
|
||||
|
||||
session.session.read_resource.assert_not_called()
|
||||
|
||||
|
||||
def test_resource_uri_allowed_supports_listed_templates_conservatively():
|
||||
session = _connected_session(
|
||||
resources=[],
|
||||
templates=[
|
||||
{
|
||||
'uri_template': 'repo://{owner}/{repo}/file/{path}',
|
||||
'name': 'repository file',
|
||||
'title': '',
|
||||
'description': '',
|
||||
'mime_type': 'text/plain',
|
||||
'icons': [],
|
||||
'annotations': {},
|
||||
'_meta': {},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
assert session.resource_uri_allowed('repo://langbot-app/LangBot/file/src/main.py') is True
|
||||
assert session.resource_uri_allowed('repo://langbot-app/LangBot/issues/1') is False
|
||||
assert session.resource_uri_allowed('https://example.com/secret') is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_loader_can_hide_synthetic_resource_tools():
|
||||
loader = MCPLoader(_app())
|
||||
session = _connected_session()
|
||||
loader.sessions = {'docs': 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)
|
||||
|
||||
assert {tool.name for tool in with_resource_tools} == {
|
||||
MCP_TOOL_LIST_RESOURCES,
|
||||
MCP_TOOL_READ_RESOURCE,
|
||||
}
|
||||
assert without_resource_tools == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
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={
|
||||
'_pipeline_bound_mcp_servers': ['srv-1'],
|
||||
'_pipeline_mcp_resource_agent_read_enabled': False,
|
||||
}
|
||||
)
|
||||
|
||||
result = await loader.invoke_tool(
|
||||
MCP_TOOL_READ_RESOURCE,
|
||||
{'server_name': 'docs', 'uri': 'file:///README.md'},
|
||||
query,
|
||||
)
|
||||
|
||||
assert result[0].text == 'Error: MCP resource agent reads are disabled.'
|
||||
session.session.read_resource.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_resource_context_for_query_uses_only_bound_attached_text_resources():
|
||||
loader = MCPLoader(_app())
|
||||
docs = _connected_session(name='docs', uuid='srv-1')
|
||||
docs.session.read_resource.return_value = mcp_types.ReadResourceResult(
|
||||
contents=[
|
||||
mcp_types.TextResourceContents(
|
||||
uri='file:///README.md',
|
||||
mimeType='text/markdown',
|
||||
text='LangBot MCP resource context',
|
||||
)
|
||||
]
|
||||
)
|
||||
other = _connected_session(name='other', uuid='srv-2')
|
||||
other.session.read_resource.return_value = mcp_types.ReadResourceResult(
|
||||
contents=[
|
||||
mcp_types.TextResourceContents(
|
||||
uri='file:///README.md',
|
||||
mimeType='text/markdown',
|
||||
text='must not be injected',
|
||||
)
|
||||
]
|
||||
)
|
||||
loader.sessions = {'docs': docs, 'other': other}
|
||||
query = SimpleNamespace(
|
||||
variables={
|
||||
'_pipeline_bound_mcp_servers': ['srv-1'],
|
||||
'_pipeline_mcp_resource_attachments': [
|
||||
{'server_uuid': 'srv-1', 'server_name': 'docs', 'uri': 'file:///README.md', 'mode': 'pinned'},
|
||||
{'server_uuid': 'srv-2', 'server_name': 'other', 'uri': 'file:///README.md', 'mode': 'pinned'},
|
||||
],
|
||||
}
|
||||
)
|
||||
|
||||
context = await loader.build_resource_context_for_query(query)
|
||||
|
||||
assert '<mcp_resource ' in context
|
||||
assert 'server="docs"' in context
|
||||
assert 'LangBot MCP resource context' in context
|
||||
assert 'must not be injected' not in context
|
||||
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()
|
||||
@@ -494,6 +494,7 @@ async def test_model_manager_init_temporary_runtime_llm_model(fake_requester_reg
|
||||
'api_keys': ['temp-key'],
|
||||
},
|
||||
'abilities': ['func_call'],
|
||||
'context_length': 128000,
|
||||
'extra_args': {'temperature': 0.5},
|
||||
}
|
||||
|
||||
@@ -501,6 +502,9 @@ async def test_model_manager_init_temporary_runtime_llm_model(fake_requester_reg
|
||||
|
||||
assert runtime_model.model_entity.uuid == 'temp-model-uuid'
|
||||
assert runtime_model.model_entity.name == 'TempModel'
|
||||
assert runtime_model.model_entity.context_length == 128000
|
||||
assert runtime_model.model_entity.extra_args == {'temperature': 0.5}
|
||||
assert 'context_length' not in runtime_model.model_entity.extra_args
|
||||
assert runtime_model.provider.provider_entity.uuid == 'temp-provider-uuid'
|
||||
assert runtime_model.provider.token_mgr.tokens == ['temp-key']
|
||||
|
||||
@@ -631,7 +635,9 @@ async def test_model_manager_reload_provider_not_found(fake_requester_registry):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_manager_load_llm_model_with_provider(fake_requester_registry, fake_persistence_data, runtime_provider):
|
||||
async def test_model_manager_load_llm_model_with_provider(
|
||||
fake_requester_registry, fake_persistence_data, runtime_provider
|
||||
):
|
||||
"""Test ModelManager.load_llm_model_with_provider creates RuntimeLLMModel."""
|
||||
model_mgr = fake_requester_registry
|
||||
|
||||
@@ -644,7 +650,9 @@ async def test_model_manager_load_llm_model_with_provider(fake_requester_registr
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_manager_load_llm_model_with_provider_from_row(fake_requester_registry, fake_persistence_data, runtime_provider):
|
||||
async def test_model_manager_load_llm_model_with_provider_from_row(
|
||||
fake_requester_registry, fake_persistence_data, runtime_provider
|
||||
):
|
||||
"""Test ModelManager.load_llm_model_with_provider handles Row objects."""
|
||||
model_mgr = fake_requester_registry
|
||||
|
||||
@@ -657,7 +665,9 @@ async def test_model_manager_load_llm_model_with_provider_from_row(fake_requeste
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_manager_load_embedding_model_with_provider(fake_requester_registry, fake_persistence_data, runtime_provider):
|
||||
async def test_model_manager_load_embedding_model_with_provider(
|
||||
fake_requester_registry, fake_persistence_data, runtime_provider
|
||||
):
|
||||
"""Test ModelManager.load_embedding_model_with_provider creates RuntimeEmbeddingModel."""
|
||||
model_mgr = fake_requester_registry
|
||||
|
||||
@@ -785,4 +795,4 @@ def test_provider_not_found_error_str():
|
||||
error = provider_errors.ProviderNotFoundError('test-provider')
|
||||
|
||||
assert str(error) == 'Provider test-provider not found'
|
||||
assert error.provider_name == 'test-provider'
|
||||
assert error.provider_name == 'test-provider'
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
@@ -16,8 +17,6 @@ from langbot.pkg.entity.persistence import model as persistence_model
|
||||
from langbot.pkg.pipeline.preproc.preproc import PreProcessor
|
||||
from langbot.pkg.provider.modelmgr import requester
|
||||
from langbot.pkg.provider.modelmgr.modelmgr import ModelManager
|
||||
from langbot.pkg.provider.modelmgr.requesters.chatcmpl import OpenAIChatCompletions
|
||||
from langbot.pkg.provider.modelmgr.requesters.modelscopechatcmpl import ModelScopeChatCompletions
|
||||
from langbot.pkg.provider.modelmgr.token import TokenManager
|
||||
from langbot.pkg.provider.runners.localagent import LocalAgentRunner
|
||||
|
||||
@@ -91,71 +90,25 @@ def test_token_manager_next_token_ignores_empty_token_list():
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_requester_initialize_uses_placeholder_api_key(monkeypatch):
|
||||
captured_kwargs = {}
|
||||
async def test_model_manager_initialize_skips_space_sync_after_timeout():
|
||||
ap = SimpleNamespace()
|
||||
ap.discover = SimpleNamespace(get_components_by_kind=Mock(return_value=[]))
|
||||
ap.instance_config = SimpleNamespace(data={'space': {'models_sync_timeout': 0.01}})
|
||||
ap.logger = Mock()
|
||||
|
||||
def fake_client(**kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
return SimpleNamespace(**kwargs)
|
||||
mgr = ModelManager(ap)
|
||||
mgr.load_models_from_db = AsyncMock()
|
||||
|
||||
monkeypatch.setattr('langbot.pkg.provider.modelmgr.requesters.chatcmpl.openai.AsyncClient', fake_client)
|
||||
monkeypatch.setattr('langbot.pkg.provider.modelmgr.requesters.chatcmpl.httpx.AsyncClient', fake_client)
|
||||
async def slow_sync():
|
||||
await asyncio.sleep(1)
|
||||
|
||||
requester_inst = OpenAIChatCompletions(ap=SimpleNamespace(), config={})
|
||||
await requester_inst.initialize()
|
||||
mgr.sync_new_models_from_space = AsyncMock(side_effect=slow_sync)
|
||||
|
||||
assert captured_kwargs['api_key'] == OpenAIChatCompletions.init_api_key
|
||||
await mgr.initialize()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_modelscope_requester_initialize_uses_placeholder_api_key(monkeypatch):
|
||||
captured_kwargs = {}
|
||||
|
||||
def fake_client(**kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
return SimpleNamespace(**kwargs)
|
||||
|
||||
monkeypatch.setattr('langbot.pkg.provider.modelmgr.requesters.modelscopechatcmpl.openai.AsyncClient', fake_client)
|
||||
monkeypatch.setattr('langbot.pkg.provider.modelmgr.requesters.modelscopechatcmpl.httpx.AsyncClient', fake_client)
|
||||
|
||||
requester_inst = ModelScopeChatCompletions(ap=SimpleNamespace(), config={})
|
||||
await requester_inst.initialize()
|
||||
|
||||
assert captured_kwargs['api_key'] == ModelScopeChatCompletions.init_api_key
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_embedding_call_overrides_placeholder_api_key():
|
||||
captured_request = {}
|
||||
|
||||
async def fake_create(**kwargs):
|
||||
captured_request['api_key'] = fake_client.api_key
|
||||
captured_request['kwargs'] = kwargs
|
||||
return SimpleNamespace(
|
||||
data=[SimpleNamespace(embedding=[0.1, 0.2])],
|
||||
usage=SimpleNamespace(prompt_tokens=3, total_tokens=3),
|
||||
)
|
||||
|
||||
fake_client = SimpleNamespace(
|
||||
api_key=OpenAIChatCompletions.init_api_key,
|
||||
embeddings=SimpleNamespace(create=fake_create),
|
||||
)
|
||||
|
||||
requester_inst = OpenAIChatCompletions(ap=SimpleNamespace(), config={})
|
||||
requester_inst.client = fake_client
|
||||
|
||||
embeddings, usage_info = await requester_inst.invoke_embedding(
|
||||
model=requester.RuntimeEmbeddingModel(
|
||||
model_entity=SimpleNamespace(name='text-embedding-3-small', extra_args={}),
|
||||
provider=SimpleNamespace(token_mgr=TokenManager('provider-uuid', [' runtime-key ', '', 'runtime-key'])),
|
||||
),
|
||||
input_text=['hello'],
|
||||
)
|
||||
|
||||
assert captured_request['api_key'] == 'runtime-key'
|
||||
assert captured_request['kwargs']['model'] == 'text-embedding-3-small'
|
||||
assert embeddings == [[0.1, 0.2]]
|
||||
assert usage_info == {'prompt_tokens': 3, 'total_tokens': 3}
|
||||
mgr.load_models_from_db.assert_awaited_once()
|
||||
mgr.sync_new_models_from_space.assert_awaited_once()
|
||||
ap.logger.warning.assert_any_call('LangBot Space model sync timed out after 0.01s, skipping startup sync.')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -169,6 +122,7 @@ async def test_updated_llm_model_is_immediately_usable_by_local_agent_pipeline()
|
||||
ap.logger = Mock()
|
||||
ap.persistence_mgr = SimpleNamespace(execute_async=AsyncMock())
|
||||
ap.tool_mgr = SimpleNamespace(get_all_tools=AsyncMock(return_value=[]))
|
||||
ap.skill_mgr = None # PreProcessor only uses skill_mgr for the local-agent skill-binding branch
|
||||
ap.plugin_connector = SimpleNamespace(
|
||||
emit_event=AsyncMock(return_value=SimpleNamespace(event=SimpleNamespace(default_prompt=[], prompt=[])))
|
||||
)
|
||||
|
||||
@@ -0,0 +1,172 @@
|
||||
"""Unit tests for provider_specific_fields round-trip in LiteLLMRequester.
|
||||
|
||||
This tests the fix for GitHub issue #1899: Gemini requires thought_signature
|
||||
to be preserved across tool call rounds for function calls to work correctly.
|
||||
"""
|
||||
|
||||
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
||||
|
||||
from langbot.pkg.provider.modelmgr.requesters.litellmchat import LiteLLMRequester
|
||||
|
||||
|
||||
def _make_requester() -> LiteLLMRequester:
|
||||
# _convert_messages and _normalize_stream_tool_calls do not touch instance config.
|
||||
return LiteLLMRequester.__new__(LiteLLMRequester)
|
||||
|
||||
|
||||
def test_convert_messages_preserves_tool_call_provider_specific_fields():
|
||||
"""Tool calls should retain provider_specific_fields through _convert_messages."""
|
||||
req = _make_requester()
|
||||
msg = provider_message.Message(
|
||||
role='assistant',
|
||||
content=None,
|
||||
tool_calls=[
|
||||
provider_message.ToolCall(
|
||||
id='call_123',
|
||||
type='function',
|
||||
function=provider_message.FunctionCall(
|
||||
name='search',
|
||||
arguments='{"query": "test"}',
|
||||
),
|
||||
provider_specific_fields={
|
||||
'thought_signature': 'c2tpcF90aG91Z2h0X3NpZ25hdHVyZQ==',
|
||||
},
|
||||
),
|
||||
],
|
||||
)
|
||||
out = req._convert_messages([msg])
|
||||
assert len(out) == 1
|
||||
assert out[0]['tool_calls'] is not None
|
||||
assert len(out[0]['tool_calls']) == 1
|
||||
|
||||
tc = out[0]['tool_calls'][0]
|
||||
assert tc['id'] == 'call_123'
|
||||
assert tc['function']['name'] == 'search'
|
||||
assert 'provider_specific_fields' in tc
|
||||
assert tc['provider_specific_fields']['thought_signature'] == 'c2tpcF90aG91Z2h0X3NpZ25hdHVyZQ=='
|
||||
|
||||
|
||||
def test_convert_messages_preserves_message_provider_specific_fields():
|
||||
"""Messages should retain provider_specific_fields through _convert_messages."""
|
||||
req = _make_requester()
|
||||
msg = provider_message.Message(
|
||||
role='assistant',
|
||||
content='Hello',
|
||||
provider_specific_fields={
|
||||
'thought_signatures': ['sig1', 'sig2'],
|
||||
},
|
||||
)
|
||||
out = req._convert_messages([msg])
|
||||
assert len(out) == 1
|
||||
assert 'provider_specific_fields' in out[0]
|
||||
assert out[0]['provider_specific_fields']['thought_signatures'] == ['sig1', 'sig2']
|
||||
|
||||
|
||||
def test_normalize_stream_tool_calls_preserves_provider_specific_fields():
|
||||
"""Streaming tool calls should retain provider_specific_fields."""
|
||||
req = _make_requester()
|
||||
tool_call_state: dict[int, dict] = {}
|
||||
|
||||
# Simulate first chunk with id and type
|
||||
raw_tool_calls_1 = [
|
||||
{
|
||||
'index': 0,
|
||||
'id': 'call_abc',
|
||||
'type': 'function',
|
||||
'function': {
|
||||
'name': 'get_weather',
|
||||
'arguments': '',
|
||||
},
|
||||
'provider_specific_fields': {
|
||||
'thought_signature': 'dGVzdF9zaWduYXR1cmU=',
|
||||
},
|
||||
},
|
||||
]
|
||||
result_1 = req._normalize_stream_tool_calls(raw_tool_calls_1, tool_call_state)
|
||||
assert result_1 is not None
|
||||
assert len(result_1) == 1
|
||||
assert result_1[0]['provider_specific_fields']['thought_signature'] == 'dGVzdF9zaWduYXR1cmU='
|
||||
|
||||
# Simulate second chunk without provider_specific_fields (should be retained from state)
|
||||
raw_tool_calls_2 = [
|
||||
{
|
||||
'index': 0,
|
||||
'function': {
|
||||
'arguments': '{"city": "Tokyo"}',
|
||||
},
|
||||
},
|
||||
]
|
||||
result_2 = req._normalize_stream_tool_calls(raw_tool_calls_2, tool_call_state)
|
||||
assert result_2 is not None
|
||||
assert len(result_2) == 1
|
||||
# Should retain the provider_specific_fields from the first chunk
|
||||
assert result_2[0]['provider_specific_fields']['thought_signature'] == 'dGVzdF9zaWduYXR1cmU='
|
||||
assert result_2[0]['function']['arguments'] == '{"city": "Tokyo"}'
|
||||
|
||||
|
||||
def test_normalize_stream_tool_calls_merges_function_level_psf():
|
||||
"""Function-level provider_specific_fields should be merged into tool-level."""
|
||||
req = _make_requester()
|
||||
tool_call_state: dict[int, dict] = {}
|
||||
|
||||
raw_tool_calls = [
|
||||
{
|
||||
'index': 0,
|
||||
'id': 'call_xyz',
|
||||
'type': 'function',
|
||||
'function': {
|
||||
'name': 'search',
|
||||
'arguments': '{}',
|
||||
'provider_specific_fields': {
|
||||
'thought_signature': 'ZnVuY19sZXZlbF9zaWc=',
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
result = req._normalize_stream_tool_calls(raw_tool_calls, tool_call_state)
|
||||
assert result is not None
|
||||
assert result[0]['provider_specific_fields']['thought_signature'] == 'ZnVuY19sZXZlbF9zaWc='
|
||||
|
||||
|
||||
def test_tool_call_roundtrip_through_message_entity():
|
||||
"""Full round-trip: LiteLLM response dict -> Message entity -> _convert_messages."""
|
||||
# Simulate what LiteLLM returns for a Gemini tool call response
|
||||
message_data = {
|
||||
'role': 'assistant',
|
||||
'content': None,
|
||||
'tool_calls': [
|
||||
{
|
||||
'id': 'call_gemini_123',
|
||||
'type': 'function',
|
||||
'function': {
|
||||
'name': 'web_search',
|
||||
'arguments': '{"query": "test"}',
|
||||
},
|
||||
'provider_specific_fields': {
|
||||
'thought_signature': 'Z2VtaW5pX3NpZ25hdHVyZQ==',
|
||||
},
|
||||
},
|
||||
],
|
||||
'provider_specific_fields': {
|
||||
'thought_signatures': ['Z2VtaW5pX3NpZ25hdHVyZQ=='],
|
||||
},
|
||||
}
|
||||
|
||||
# Parse into Message entity (this is what invoke_llm does)
|
||||
msg = provider_message.Message(**message_data)
|
||||
|
||||
# Verify the entity has the fields
|
||||
assert msg.tool_calls is not None
|
||||
assert len(msg.tool_calls) == 1
|
||||
assert msg.tool_calls[0].provider_specific_fields is not None
|
||||
assert msg.tool_calls[0].provider_specific_fields['thought_signature'] == 'Z2VtaW5pX3NpZ25hdHVyZQ=='
|
||||
assert msg.provider_specific_fields is not None
|
||||
assert msg.provider_specific_fields['thought_signatures'] == ['Z2VtaW5pX3NpZ25hdHVyZQ==']
|
||||
|
||||
# Convert back to dict for LiteLLM (this is what _convert_messages does)
|
||||
req = _make_requester()
|
||||
out = req._convert_messages([msg])
|
||||
|
||||
# Verify the fields are preserved in the output
|
||||
assert out[0]['tool_calls'][0]['provider_specific_fields']['thought_signature'] == 'Z2VtaW5pX3NpZ25hdHVyZQ=='
|
||||
assert out[0]['provider_specific_fields']['thought_signatures'] == ['Z2VtaW5pX3NpZ25hdHVyZQ==']
|
||||
@@ -43,6 +43,7 @@ class TestableRequester(requester.ProviderAPIRequester):
|
||||
remove_think=False,
|
||||
):
|
||||
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
||||
|
||||
return provider_message.Message(
|
||||
role='assistant',
|
||||
content=[provider_message.ContentElement(type='text', text='Testable response')],
|
||||
@@ -289,7 +290,9 @@ async def test_runtime_provider_invoke_llm_delegates(runtime_provider, runtime_l
|
||||
current_stage_name=None,
|
||||
)
|
||||
|
||||
messages = [provider_message.Message(role='user', content=[provider_message.ContentElement(type='text', text='Hello')])]
|
||||
messages = [
|
||||
provider_message.Message(role='user', content=[provider_message.ContentElement(type='text', text='Hello')])
|
||||
]
|
||||
|
||||
result = await provider.invoke_llm(query, runtime_llm_model, messages)
|
||||
|
||||
@@ -330,7 +333,9 @@ async def test_runtime_provider_invoke_llm_stream_yields_chunks(runtime_provider
|
||||
current_stage_name=None,
|
||||
)
|
||||
|
||||
messages = [provider_message.Message(role='user', content=[provider_message.ContentElement(type='text', text='Hello')])]
|
||||
messages = [
|
||||
provider_message.Message(role='user', content=[provider_message.ContentElement(type='text', text='Hello')])
|
||||
]
|
||||
|
||||
chunks = []
|
||||
async for chunk in provider.invoke_llm_stream(query, runtime_llm_model, messages):
|
||||
@@ -576,7 +581,9 @@ async def test_runtime_provider_invoke_llm_propagates_error(mock_app_for_modelmg
|
||||
current_stage_name=None,
|
||||
)
|
||||
|
||||
messages = [provider_message.Message(role='user', content=[provider_message.ContentElement(type='text', text='Hello')])]
|
||||
messages = [
|
||||
provider_message.Message(role='user', content=[provider_message.ContentElement(type='text', text='Hello')])
|
||||
]
|
||||
|
||||
with pytest.raises(RequesterError):
|
||||
await provider.invoke_llm(query, model, messages)
|
||||
|
||||
@@ -5,6 +5,7 @@ Tests cover:
|
||||
- Conversation creation with prompts
|
||||
- Session concurrency semaphore
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
@@ -60,11 +61,7 @@ class TestSessionManagerGetSession:
|
||||
"""Create mock app with instance config."""
|
||||
mock_app = Mock()
|
||||
mock_app.instance_config = Mock()
|
||||
mock_app.instance_config.data = {
|
||||
'concurrency': {
|
||||
'session': 5
|
||||
}
|
||||
}
|
||||
mock_app.instance_config.data = {'concurrency': {'session': 5}}
|
||||
return mock_app
|
||||
|
||||
@pytest.fixture
|
||||
@@ -173,11 +170,7 @@ class TestSessionManagerGetConversation:
|
||||
"""Create mock app with instance config."""
|
||||
mock_app = Mock()
|
||||
mock_app.instance_config = Mock()
|
||||
mock_app.instance_config.data = {
|
||||
'concurrency': {
|
||||
'session': 5
|
||||
}
|
||||
}
|
||||
mock_app.instance_config.data = {'concurrency': {'session': 5}}
|
||||
return mock_app
|
||||
|
||||
@pytest.fixture
|
||||
@@ -201,17 +194,13 @@ class TestSessionManagerGetConversation:
|
||||
return query
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_creates_conversation_with_prompt(
|
||||
self, mock_app_with_config, sample_query, sample_session
|
||||
):
|
||||
async def test_creates_conversation_with_prompt(self, mock_app_with_config, sample_query, sample_session):
|
||||
"""Test that get_conversation creates conversation with prompt."""
|
||||
sessionmgr = get_session_module()
|
||||
|
||||
manager = sessionmgr.SessionManager(mock_app_with_config)
|
||||
|
||||
prompt_config = [
|
||||
{'role': 'system', 'content': 'You are a helpful assistant.'}
|
||||
]
|
||||
prompt_config = [{'role': 'system', 'content': 'You are a helpful assistant.'}]
|
||||
pipeline_uuid = 'pipeline-123'
|
||||
bot_uuid = 'bot-123'
|
||||
|
||||
@@ -234,21 +223,15 @@ class TestSessionManagerGetConversation:
|
||||
|
||||
manager = sessionmgr.SessionManager(mock_app_with_config)
|
||||
|
||||
prompt_config = [
|
||||
{'role': 'system', 'content': 'You are a helpful assistant.'}
|
||||
]
|
||||
prompt_config = [{'role': 'system', 'content': 'You are a helpful assistant.'}]
|
||||
pipeline_uuid = 'pipeline-123'
|
||||
bot_uuid = 'bot-123'
|
||||
|
||||
# First call creates conversation
|
||||
conv1 = await manager.get_conversation(
|
||||
sample_query, sample_session, prompt_config, pipeline_uuid, bot_uuid
|
||||
)
|
||||
conv1 = await manager.get_conversation(sample_query, sample_session, prompt_config, pipeline_uuid, bot_uuid)
|
||||
|
||||
# Second call with same pipeline should return same conversation
|
||||
conv2 = await manager.get_conversation(
|
||||
sample_query, sample_session, prompt_config, pipeline_uuid, bot_uuid
|
||||
)
|
||||
conv2 = await manager.get_conversation(sample_query, sample_session, prompt_config, pipeline_uuid, bot_uuid)
|
||||
|
||||
assert conv1 is conv2
|
||||
assert len(sample_session.conversations) == 1
|
||||
@@ -262,36 +245,26 @@ class TestSessionManagerGetConversation:
|
||||
|
||||
manager = sessionmgr.SessionManager(mock_app_with_config)
|
||||
|
||||
prompt_config = [
|
||||
{'role': 'system', 'content': 'You are a helpful assistant.'}
|
||||
]
|
||||
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', 'bot-1')
|
||||
|
||||
# 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', 'bot-2')
|
||||
|
||||
assert conv1 is not conv2
|
||||
assert len(sample_session.conversations) == 2
|
||||
assert sample_session.using_conversation is conv2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conversation_has_empty_messages(
|
||||
self, mock_app_with_config, sample_query, sample_session
|
||||
):
|
||||
async def test_conversation_has_empty_messages(self, mock_app_with_config, sample_query, sample_session):
|
||||
"""Test that created conversation has empty messages list."""
|
||||
sessionmgr = get_session_module()
|
||||
|
||||
manager = sessionmgr.SessionManager(mock_app_with_config)
|
||||
|
||||
prompt_config = [
|
||||
{'role': 'system', 'content': 'You are a helpful assistant.'}
|
||||
]
|
||||
prompt_config = [{'role': 'system', 'content': 'You are a helpful assistant.'}]
|
||||
|
||||
conversation = await manager.get_conversation(
|
||||
sample_query, sample_session, prompt_config, 'pipeline-123', 'bot-123'
|
||||
@@ -300,22 +273,17 @@ class TestSessionManagerGetConversation:
|
||||
assert conversation.messages == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prompt_messages_from_config(
|
||||
self, mock_app_with_config, sample_query, sample_session
|
||||
):
|
||||
async def test_prompt_messages_from_config(self, mock_app_with_config, sample_query, sample_session):
|
||||
"""Test that prompt messages are created from prompt_config."""
|
||||
sessionmgr = get_session_module()
|
||||
|
||||
manager = sessionmgr.SessionManager(mock_app_with_config)
|
||||
|
||||
prompt_config = [
|
||||
{'role': 'system', 'content': 'System message'},
|
||||
{'role': 'user', 'content': 'User message'}
|
||||
]
|
||||
prompt_config = [{'role': 'system', 'content': 'System message'}, {'role': 'user', 'content': 'User message'}]
|
||||
|
||||
conversation = await manager.get_conversation(
|
||||
sample_query, sample_session, prompt_config, 'pipeline-123', 'bot-123'
|
||||
)
|
||||
|
||||
assert conversation.prompt.name == 'default'
|
||||
assert len(conversation.prompt.messages) == 2
|
||||
assert len(conversation.prompt.messages) == 2
|
||||
|
||||
@@ -0,0 +1,539 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def _make_ap(logger=None):
|
||||
ap = SimpleNamespace()
|
||||
ap.logger = logger or Mock()
|
||||
ap.persistence_mgr = Mock()
|
||||
ap.persistence_mgr.execute_async = AsyncMock(return_value=Mock(all=Mock(return_value=[])))
|
||||
ap.persistence_mgr.serialize_model = Mock(side_effect=lambda cls, row: row)
|
||||
return ap
|
||||
|
||||
|
||||
def _make_skill_data(
|
||||
name='test-skill',
|
||||
instructions='Do something',
|
||||
package_root='',
|
||||
entry_file='SKILL.md',
|
||||
**kwargs,
|
||||
):
|
||||
return {
|
||||
'name': name,
|
||||
'display_name': kwargs.pop('display_name', name),
|
||||
'description': kwargs.pop('description', f'Description of {name}'),
|
||||
'instructions': instructions,
|
||||
'package_root': package_root,
|
||||
'entry_file': entry_file,
|
||||
**kwargs,
|
||||
}
|
||||
|
||||
|
||||
class TestSkillManagerCache:
|
||||
"""The Box runtime is the only source of truth — SkillManager just holds
|
||||
an in-memory cache populated by ``reload_skills``. There is no local
|
||||
filesystem reader anymore."""
|
||||
|
||||
def test_refresh_skill_from_disk_reports_cache_presence(self):
|
||||
"""Box is the only source of truth for skill content. refresh_skill_from_disk
|
||||
now just reports whether the skill is still in the in-memory cache —
|
||||
the actual content refresh is driven by SkillService awaiting
|
||||
``reload_skills`` after every Box mutation."""
|
||||
from langbot.pkg.skill.manager import SkillManager
|
||||
|
||||
ap = _make_ap()
|
||||
mgr = SkillManager(ap)
|
||||
|
||||
# Empty cache → returns False
|
||||
assert mgr.refresh_skill_from_disk('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
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reload_skills_drops_box_skills_with_missing_package_root(self):
|
||||
"""When LangBot shares a filesystem with Box (local stdio mode) and Box
|
||||
reports a skill whose package_root is gone from that shared filesystem,
|
||||
the cache must drop it instead of keeping a stale entry that would later
|
||||
produce a bad mount."""
|
||||
from langbot.pkg.skill.manager import SkillManager
|
||||
|
||||
with tempfile.TemporaryDirectory() as live_dir:
|
||||
ghost_dir = os.path.join(live_dir, '_does_not_exist')
|
||||
box_service = SimpleNamespace(
|
||||
available=True,
|
||||
shares_filesystem_with_box=True,
|
||||
list_skills=AsyncMock(
|
||||
return_value=[
|
||||
_make_skill_data(name='alive', package_root=live_dir),
|
||||
_make_skill_data(name='ghost', package_root=ghost_dir),
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
ap = _make_ap()
|
||||
ap.box_service = box_service
|
||||
mgr = SkillManager(ap)
|
||||
|
||||
await mgr.reload_skills()
|
||||
|
||||
assert list(mgr.skills) == ['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)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reload_skills_trusts_box_paths_when_filesystem_not_shared(self):
|
||||
"""In separated deployments (Docker Compose, k8s sidecar,
|
||||
--standalone-box, remote endpoint) the package_root reported by Box
|
||||
lives on the Box runtime's filesystem and is not resolvable on the
|
||||
LangBot side. The cache must keep every Box-reported skill rather than
|
||||
dropping them all via a local isdir() check."""
|
||||
from langbot.pkg.skill.manager import SkillManager
|
||||
|
||||
box_service = SimpleNamespace(
|
||||
available=True,
|
||||
shares_filesystem_with_box=False,
|
||||
list_skills=AsyncMock(
|
||||
return_value=[
|
||||
_make_skill_data(name='alpha', package_root='/box/skills/alpha'),
|
||||
_make_skill_data(name='beta', package_root='/box/skills/beta'),
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
ap = _make_ap()
|
||||
ap.box_service = box_service
|
||||
mgr = SkillManager(ap)
|
||||
|
||||
await mgr.reload_skills()
|
||||
|
||||
assert sorted(mgr.skills) == ['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)
|
||||
|
||||
|
||||
class TestSkillActivationHelper:
|
||||
"""Skill activation is now Tool-Call based.
|
||||
|
||||
The legacy text-marker mechanism (``[ACTIVATE_SKILL: x]`` detection,
|
||||
``build_activation_prompt_for_skills``, ``remove_activation_marker``,
|
||||
``prepare_skill_activation``) has been removed. Activation now goes
|
||||
through ``skill.activation.register_activated_skill``, invoked by the
|
||||
``activate`` Tool Call.
|
||||
"""
|
||||
|
||||
def test_register_activated_skill_records_known_skill(self):
|
||||
from langbot.pkg.skill.activation import register_activated_skill
|
||||
from langbot.pkg.provider.tools.loaders.skill import ACTIVATED_SKILLS_KEY
|
||||
from langbot.pkg.skill.manager import SkillManager
|
||||
|
||||
ap = _make_ap()
|
||||
mgr = SkillManager(ap)
|
||||
mgr.skills = {
|
||||
'primary': _make_skill_data(name='primary', instructions='Primary instructions'),
|
||||
}
|
||||
ap.skill_mgr = mgr
|
||||
|
||||
query = SimpleNamespace(variables={})
|
||||
|
||||
assert register_activated_skill(ap, query, 'primary') is True
|
||||
assert set(query.variables[ACTIVATED_SKILLS_KEY].keys()) == {'primary'}
|
||||
assert query.variables[ACTIVATED_SKILLS_KEY]['primary']['name'] == 'primary'
|
||||
|
||||
def test_register_activated_skill_rejects_unknown_skill(self):
|
||||
from langbot.pkg.skill.activation import register_activated_skill
|
||||
from langbot.pkg.provider.tools.loaders.skill import ACTIVATED_SKILLS_KEY
|
||||
from langbot.pkg.skill.manager import SkillManager
|
||||
|
||||
ap = _make_ap()
|
||||
mgr = SkillManager(ap)
|
||||
mgr.skills = {'primary': _make_skill_data(name='primary')}
|
||||
ap.skill_mgr = mgr
|
||||
|
||||
query = SimpleNamespace(variables={})
|
||||
|
||||
assert register_activated_skill(ap, query, 'missing') is False
|
||||
assert ACTIVATED_SKILLS_KEY not in query.variables
|
||||
|
||||
def test_register_activated_skill_without_skill_manager_returns_false(self):
|
||||
from langbot.pkg.skill.activation import register_activated_skill
|
||||
|
||||
ap = _make_ap() # no skill_mgr attribute
|
||||
query = SimpleNamespace(variables={})
|
||||
|
||||
assert register_activated_skill(ap, query, 'primary') is False
|
||||
|
||||
|
||||
class TestSkillPathHelpers:
|
||||
def test_get_visible_skills_filters_by_bound_names(self):
|
||||
from langbot.pkg.provider.tools.loaders.skill import PIPELINE_BOUND_SKILLS_KEY, get_visible_skills
|
||||
|
||||
ap = _make_ap()
|
||||
ap.skill_mgr = SimpleNamespace(
|
||||
skills={
|
||||
'visible': _make_skill_data(name='visible'),
|
||||
'hidden': _make_skill_data(name='hidden'),
|
||||
}
|
||||
)
|
||||
query = SimpleNamespace(variables={PIPELINE_BOUND_SKILLS_KEY: ['visible']})
|
||||
|
||||
result = get_visible_skills(ap, query)
|
||||
|
||||
assert list(result.keys()) == ['visible']
|
||||
|
||||
def test_restore_activated_skills_uses_caller_provided_names_and_visibility(self):
|
||||
from langbot.pkg.provider.tools.loaders.skill import (
|
||||
ACTIVATED_SKILLS_KEY,
|
||||
PIPELINE_BOUND_SKILLS_KEY,
|
||||
get_activated_skill_names,
|
||||
restore_activated_skills,
|
||||
)
|
||||
|
||||
ap = _make_ap()
|
||||
ap.skill_mgr = SimpleNamespace(
|
||||
skills={
|
||||
'visible': _make_skill_data(name='visible'),
|
||||
'hidden': _make_skill_data(name='hidden'),
|
||||
}
|
||||
)
|
||||
query = SimpleNamespace(variables={PIPELINE_BOUND_SKILLS_KEY: ['visible']})
|
||||
|
||||
restored = restore_activated_skills(ap, query, ['visible', 'hidden', 'visible', ''])
|
||||
|
||||
assert restored == ['visible']
|
||||
assert list(query.variables[ACTIVATED_SKILLS_KEY].keys()) == ['visible']
|
||||
assert get_activated_skill_names(query) == ['visible']
|
||||
|
||||
def test_resolve_virtual_skill_path_allows_visible_skill_reads(self):
|
||||
from langbot.pkg.provider.tools.loaders.skill import (
|
||||
PIPELINE_BOUND_SKILLS_KEY,
|
||||
resolve_virtual_skill_path,
|
||||
)
|
||||
|
||||
ap = _make_ap()
|
||||
ap.skill_mgr = SimpleNamespace(skills={'demo': _make_skill_data(name='demo')})
|
||||
query = SimpleNamespace(variables={PIPELINE_BOUND_SKILLS_KEY: ['demo']})
|
||||
|
||||
skill, rewritten = resolve_virtual_skill_path(
|
||||
ap,
|
||||
query,
|
||||
'/workspace/.skills/demo/SKILL.md',
|
||||
include_visible=True,
|
||||
include_activated=False,
|
||||
)
|
||||
|
||||
assert skill['name'] == 'demo'
|
||||
assert rewritten == '/workspace/SKILL.md'
|
||||
|
||||
def test_build_skill_session_id_uses_name_based_identifier(self):
|
||||
from langbot.pkg.provider.tools.loaders.skill import build_skill_session_id
|
||||
|
||||
with_launcher = build_skill_session_id(
|
||||
{'name': 'writer'},
|
||||
SimpleNamespace(query_id=42, launcher_type='person', launcher_id='123'),
|
||||
)
|
||||
fallback = build_skill_session_id({'name': 'writer'}, SimpleNamespace(query_id=99))
|
||||
|
||||
assert with_launcher == 'skill-person_123-writer'
|
||||
assert fallback == 'skill-99-writer'
|
||||
|
||||
def test_should_prepare_skill_python_env_detects_manifests_and_venv(self):
|
||||
from langbot.pkg.provider.tools.loaders.skill import should_prepare_skill_python_env
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
assert should_prepare_skill_python_env(tmpdir) is False
|
||||
|
||||
with open(os.path.join(tmpdir, 'requirements.txt'), 'w', encoding='utf-8') as f:
|
||||
f.write('requests==2.32.0\n')
|
||||
assert should_prepare_skill_python_env(tmpdir) is True
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
os.makedirs(os.path.join(tmpdir, '.venv'))
|
||||
assert should_prepare_skill_python_env(tmpdir) is True
|
||||
|
||||
def test_wrap_skill_command_with_python_env_bootstraps_then_runs_command(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')
|
||||
|
||||
assert '_LB_SYSTEM_PYTHON="$(command -v python3 || command -v python || true)"' in command
|
||||
assert '"$_LB_SYSTEM_PYTHON" -m venv "$_LB_VENV_DIR"' in command
|
||||
assert 'export VIRTUAL_ENV="$_LB_VENV_DIR"' in command
|
||||
assert command.rstrip().endswith('python scripts/run.py')
|
||||
|
||||
|
||||
class TestSkillToolLoader:
|
||||
"""The skill tool surface is now just ``activate`` + ``register_skill``.
|
||||
|
||||
The legacy CRUD authoring tools (create/list/get/update/delete/
|
||||
import_skill_from_directory/reload_skills) were removed; skill CRUD is
|
||||
handled by SkillService via the HTTP API / web UI instead.
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_activate_returns_instructions_and_registers_skill(self):
|
||||
from langbot.pkg.provider.tools.loaders.skill_authoring import (
|
||||
ACTIVATE_SKILL_TOOL_NAME,
|
||||
SkillToolLoader,
|
||||
)
|
||||
from langbot.pkg.provider.tools.loaders.skill import ACTIVATED_SKILLS_KEY
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
loader = SkillToolLoader(ap)
|
||||
query = SimpleNamespace(variables={})
|
||||
|
||||
result = await loader.invoke_tool(ACTIVATE_SKILL_TOOL_NAME, {'skill_name': 'demo'}, query)
|
||||
|
||||
assert result['activated'] is True
|
||||
assert result['skill_name'] == 'demo'
|
||||
assert result['mount_path'] == '/workspace/.skills/demo'
|
||||
assert result['activated_skill_names'] == ['demo']
|
||||
assert 'Step 1' in result['content']
|
||||
assert set(query.variables[ACTIVATED_SKILLS_KEY].keys()) == {'demo'}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_activate_unknown_skill_raises(self):
|
||||
from langbot.pkg.provider.tools.loaders.skill_authoring import (
|
||||
ACTIVATE_SKILL_TOOL_NAME,
|
||||
SkillToolLoader,
|
||||
)
|
||||
|
||||
ap = _make_ap()
|
||||
ap.skill_mgr = SimpleNamespace(
|
||||
skills={'demo': _make_skill_data(name='demo')},
|
||||
get_skill_by_name=lambda name: None,
|
||||
)
|
||||
|
||||
loader = SkillToolLoader(ap)
|
||||
|
||||
with pytest.raises(ValueError, match='not found'):
|
||||
await loader.invoke_tool(
|
||||
ACTIVATE_SKILL_TOOL_NAME,
|
||||
{'skill_name': 'ghost'},
|
||||
SimpleNamespace(variables={}),
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_skill_scans_directory_and_creates_skill(self):
|
||||
from langbot.pkg.provider.tools.loaders.skill_authoring import (
|
||||
REGISTER_SKILL_TOOL_NAME,
|
||||
SkillToolLoader,
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
repo_dir = os.path.join(tmpdir, 'repo')
|
||||
os.makedirs(repo_dir)
|
||||
|
||||
ap = _make_ap()
|
||||
ap.box_service = SimpleNamespace(default_workspace=tmpdir, available=True)
|
||||
ap.skill_service = SimpleNamespace(
|
||||
scan_directory_async=AsyncMock(
|
||||
return_value={
|
||||
'name': 'cloned-skill',
|
||||
'display_name': 'Cloned Skill',
|
||||
'description': 'Imported from clone',
|
||||
'instructions': 'Do work',
|
||||
}
|
||||
),
|
||||
create_skill=AsyncMock(
|
||||
return_value=_make_skill_data(name='cloned-skill', package_root=os.path.realpath(repo_dir))
|
||||
),
|
||||
)
|
||||
|
||||
loader = SkillToolLoader(ap)
|
||||
result = await loader.invoke_tool(
|
||||
REGISTER_SKILL_TOOL_NAME,
|
||||
{'path': '/workspace/repo'},
|
||||
SimpleNamespace(),
|
||||
)
|
||||
|
||||
ap.skill_service.scan_directory_async.assert_awaited_once_with(os.path.realpath(repo_dir))
|
||||
ap.skill_service.create_skill.assert_awaited_once_with(
|
||||
{
|
||||
'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'
|
||||
assert result['source_path'] == '/workspace/repo'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_skill_rejects_workspace_escape(self):
|
||||
from langbot.pkg.provider.tools.loaders.skill_authoring import (
|
||||
REGISTER_SKILL_TOOL_NAME,
|
||||
SkillToolLoader,
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
ap = _make_ap()
|
||||
ap.box_service = SimpleNamespace(default_workspace=tmpdir, available=True)
|
||||
ap.skill_service = SimpleNamespace(scan_directory_async=AsyncMock(), create_skill=AsyncMock())
|
||||
|
||||
loader = SkillToolLoader(ap)
|
||||
|
||||
with pytest.raises(ValueError, match='escapes the workspace boundary'):
|
||||
await loader.invoke_tool(
|
||||
REGISTER_SKILL_TOOL_NAME,
|
||||
{'path': '/workspace/../../etc'},
|
||||
SimpleNamespace(),
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_skill_requires_skill_service(self):
|
||||
from langbot.pkg.provider.tools.loaders.skill_authoring import (
|
||||
REGISTER_SKILL_TOOL_NAME,
|
||||
SkillToolLoader,
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
ap = _make_ap() # no skill_service attribute
|
||||
ap.box_service = SimpleNamespace(default_workspace=tmpdir, available=True)
|
||||
|
||||
loader = SkillToolLoader(ap)
|
||||
|
||||
with pytest.raises(ValueError, match='Skill service not available'):
|
||||
await loader.invoke_tool(
|
||||
REGISTER_SKILL_TOOL_NAME,
|
||||
{'path': '/workspace/foo'},
|
||||
SimpleNamespace(),
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tools_hidden_when_sandbox_backend_unavailable(self):
|
||||
from langbot.pkg.provider.tools.loaders.skill_authoring import SkillToolLoader
|
||||
|
||||
ap = _make_ap()
|
||||
ap.skill_mgr = SimpleNamespace(skills={})
|
||||
ap.box_service = SimpleNamespace(
|
||||
available=True,
|
||||
get_status=AsyncMock(return_value={'backend': {'available': False}}),
|
||||
)
|
||||
|
||||
loader = SkillToolLoader(ap)
|
||||
await loader.initialize()
|
||||
|
||||
assert await loader.get_tools() == []
|
||||
assert await loader.has_tool('activate') is False
|
||||
assert await loader.has_tool('register_skill') is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tools_exposed_when_sandbox_backend_available(self):
|
||||
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.box_service = SimpleNamespace(
|
||||
available=True,
|
||||
get_status=AsyncMock(return_value={'backend': {'available': True}}),
|
||||
)
|
||||
|
||||
loader = SkillToolLoader(ap)
|
||||
await loader.initialize()
|
||||
|
||||
tools = await loader.get_tools()
|
||||
|
||||
assert sorted(tool.name for tool in tools) == ['activate', 'register_skill']
|
||||
assert await loader.has_tool('activate') is True
|
||||
assert await loader.has_tool('register_skill') is True
|
||||
|
||||
|
||||
class TestNativeToolLoaderSkillPaths:
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_visible_skill_file(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:
|
||||
skill_md = os.path.join(tmpdir, 'SKILL.md')
|
||||
with open(skill_md, 'w', encoding='utf-8') as f:
|
||||
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)})
|
||||
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']}),
|
||||
)
|
||||
|
||||
assert result['ok'] is True
|
||||
assert result['content'] == 'demo instructions'
|
||||
assert result['truncated'] is False
|
||||
|
||||
@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
|
||||
from langbot.pkg.provider.tools.loaders.skill import register_activated_skill
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
ap = _make_ap()
|
||||
ap.box_service = SimpleNamespace(
|
||||
available=True,
|
||||
default_workspace=tmpdir,
|
||||
execute_tool=AsyncMock(return_value={'ok': True}),
|
||||
)
|
||||
ap.skill_mgr = SimpleNamespace(refresh_skill_from_disk=Mock())
|
||||
loader = NativeToolLoader(ap)
|
||||
|
||||
query = SimpleNamespace(query_id='q1', launcher_type='person', launcher_id='123', variables={})
|
||||
register_activated_skill(query, _make_skill_data(name='demo', package_root=tmpdir))
|
||||
|
||||
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]
|
||||
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')
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_requires_skill_activation(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:
|
||||
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)})
|
||||
loader = NativeToolLoader(ap)
|
||||
|
||||
query = SimpleNamespace(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(
|
||||
'write',
|
||||
{'path': '/workspace/.skills/demo/notes.txt', 'content': 'hi'},
|
||||
query,
|
||||
)
|
||||
@@ -1,9 +1,10 @@
|
||||
"""Unit tests for ToolManager.
|
||||
|
||||
Tests cover:
|
||||
- Tool schema generation for OpenAI and Anthropic
|
||||
- Tool schema generation for OpenAI/LiteLLM
|
||||
- Tool execution dispatch
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
@@ -52,11 +53,12 @@ class TestToolManagerSchemaGeneration:
|
||||
@pytest.fixture
|
||||
def sample_tools(self):
|
||||
"""Create sample LLMTool list for testing."""
|
||||
|
||||
def dummy_weather_func(**kwargs):
|
||||
return "weather result"
|
||||
return 'weather result'
|
||||
|
||||
def dummy_calc_func(**kwargs):
|
||||
return "calc result"
|
||||
return 'calc result'
|
||||
|
||||
tools = [
|
||||
resource_tool.LLMTool(
|
||||
@@ -65,15 +67,10 @@ class TestToolManagerSchemaGeneration:
|
||||
description='Get current weather for a location',
|
||||
parameters={
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'location': {
|
||||
'type': 'string',
|
||||
'description': 'City name'
|
||||
}
|
||||
},
|
||||
'required': ['location']
|
||||
'properties': {'location': {'type': 'string', 'description': 'City name'}},
|
||||
'required': ['location'],
|
||||
},
|
||||
func=dummy_weather_func
|
||||
func=dummy_weather_func,
|
||||
),
|
||||
resource_tool.LLMTool(
|
||||
name='calculate',
|
||||
@@ -81,15 +78,10 @@ class TestToolManagerSchemaGeneration:
|
||||
description='Perform a calculation',
|
||||
parameters={
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'expression': {
|
||||
'type': 'string',
|
||||
'description': 'Math expression'
|
||||
}
|
||||
},
|
||||
'required': ['expression']
|
||||
'properties': {'expression': {'type': 'string', 'description': 'Math expression'}},
|
||||
'required': ['expression'],
|
||||
},
|
||||
func=dummy_calc_func
|
||||
func=dummy_calc_func,
|
||||
),
|
||||
]
|
||||
return tools
|
||||
@@ -117,28 +109,6 @@ class TestToolManagerSchemaGeneration:
|
||||
assert tool2['type'] == 'function'
|
||||
assert tool2['function']['name'] == 'calculate'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_tools_for_anthropic(self, mock_app, sample_tools):
|
||||
"""Test that generate_tools_for_anthropic produces correct schema."""
|
||||
toolmgr = get_toolmgr_module()
|
||||
|
||||
manager = toolmgr.ToolManager(mock_app)
|
||||
result = await manager.generate_tools_for_anthropic(sample_tools)
|
||||
|
||||
assert len(result) == 2
|
||||
|
||||
# Verify first tool schema (Anthropic format)
|
||||
tool1 = result[0]
|
||||
assert tool1['name'] == 'get_weather'
|
||||
assert tool1['description'] == 'Get current weather for a location'
|
||||
assert 'input_schema' in tool1
|
||||
assert tool1['input_schema']['type'] == 'object'
|
||||
|
||||
# Verify second tool schema
|
||||
tool2 = result[1]
|
||||
assert tool2['name'] == 'calculate'
|
||||
assert 'input_schema' in tool2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_tools_empty_list(self, mock_app):
|
||||
"""Test that generating tools from empty list returns empty list."""
|
||||
@@ -149,9 +119,6 @@ class TestToolManagerSchemaGeneration:
|
||||
openai_result = await manager.generate_tools_for_openai([])
|
||||
assert openai_result == []
|
||||
|
||||
anthropic_result = await manager.generate_tools_for_anthropic([])
|
||||
assert anthropic_result == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_schema_fields_complete(self, mock_app, sample_tools):
|
||||
"""Test that OpenAI schema includes all required fields."""
|
||||
@@ -169,45 +136,54 @@ class TestToolManagerSchemaGeneration:
|
||||
assert 'description' in func
|
||||
assert 'parameters' in func
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_schema_fields_complete(self, mock_app, sample_tools):
|
||||
"""Test that Anthropic schema includes all required fields."""
|
||||
toolmgr = get_toolmgr_module()
|
||||
|
||||
manager = toolmgr.ToolManager(mock_app)
|
||||
result = await manager.generate_tools_for_anthropic(sample_tools)
|
||||
|
||||
for tool_schema in result:
|
||||
assert 'name' in tool_schema
|
||||
assert 'description' in tool_schema
|
||||
assert 'input_schema' in tool_schema
|
||||
|
||||
|
||||
class TestToolManagerExecuteFuncCall:
|
||||
"""Tests for execute_func_call method."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_app_with_loaders(self):
|
||||
"""Create mock app with mock tool loaders."""
|
||||
"""Create mock app with mock tool loaders.
|
||||
|
||||
Returns (app, plugin_loader, mcp_loader). The native and skill loaders
|
||||
are attached directly to the app for tests that don't need to assert
|
||||
against them — they all default to ``has_tool == False`` so the
|
||||
execute_func_call probe falls through to the plugin/mcp pair.
|
||||
"""
|
||||
mock_app = Mock()
|
||||
mock_app.logger = Mock()
|
||||
|
||||
def _make_inert_loader():
|
||||
loader = Mock()
|
||||
loader.has_tool = AsyncMock(return_value=False)
|
||||
loader.invoke_tool = AsyncMock(return_value=None)
|
||||
loader.initialize = AsyncMock()
|
||||
loader.shutdown = AsyncMock()
|
||||
return loader
|
||||
|
||||
# Create mock plugin loader
|
||||
mock_plugin_loader = Mock()
|
||||
mock_plugin_loader.has_tool = AsyncMock(return_value=False)
|
||||
mock_plugin_loader = _make_inert_loader()
|
||||
mock_plugin_loader.invoke_tool = AsyncMock(return_value='plugin_result')
|
||||
mock_plugin_loader.initialize = AsyncMock()
|
||||
mock_plugin_loader.shutdown = AsyncMock()
|
||||
|
||||
# Create mock MCP loader
|
||||
mock_mcp_loader = Mock()
|
||||
mock_mcp_loader.has_tool = AsyncMock(return_value=False)
|
||||
mock_mcp_loader = _make_inert_loader()
|
||||
mock_mcp_loader.invoke_tool = AsyncMock(return_value='mcp_result')
|
||||
mock_mcp_loader.initialize = AsyncMock()
|
||||
mock_mcp_loader.shutdown = AsyncMock()
|
||||
|
||||
# Stash inert native/skill loaders so the ToolManager probe order
|
||||
# (native → plugin → mcp → skill) doesn't AttributeError. Tests that
|
||||
# need to override these can replace the attributes on the manager.
|
||||
mock_app._inert_native_loader = _make_inert_loader()
|
||||
mock_app._inert_skill_loader = _make_inert_loader()
|
||||
|
||||
return mock_app, mock_plugin_loader, mock_mcp_loader
|
||||
|
||||
@staticmethod
|
||||
def _wire_loaders(manager, mock_app, plugin_loader, mcp_loader):
|
||||
"""Attach all four loaders (native + plugin + mcp + skill) to manager."""
|
||||
manager.native_tool_loader = mock_app._inert_native_loader
|
||||
manager.plugin_tool_loader = plugin_loader
|
||||
manager.mcp_tool_loader = mcp_loader
|
||||
manager.skill_tool_loader = mock_app._inert_skill_loader
|
||||
|
||||
@pytest.fixture
|
||||
def sample_query(self):
|
||||
"""Create sample query for testing."""
|
||||
@@ -215,9 +191,7 @@ class TestToolManagerExecuteFuncCall:
|
||||
return query
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_calls_plugin_loader_when_has_tool(
|
||||
self, mock_app_with_loaders, sample_query
|
||||
):
|
||||
async def test_execute_calls_plugin_loader_when_has_tool(self, mock_app_with_loaders, sample_query):
|
||||
"""Test that execute_func_call uses plugin loader when tool exists there."""
|
||||
toolmgr = get_toolmgr_module()
|
||||
|
||||
@@ -225,26 +199,17 @@ class TestToolManagerExecuteFuncCall:
|
||||
mock_plugin_loader.has_tool = AsyncMock(return_value=True)
|
||||
|
||||
manager = toolmgr.ToolManager(mock_app)
|
||||
manager.plugin_tool_loader = mock_plugin_loader
|
||||
manager.mcp_tool_loader = mock_mcp_loader
|
||||
self._wire_loaders(manager, mock_app, mock_plugin_loader, mock_mcp_loader)
|
||||
|
||||
result = await manager.execute_func_call(
|
||||
'test_tool',
|
||||
{'param': 'value'},
|
||||
sample_query
|
||||
)
|
||||
result = await manager.execute_func_call('test_tool', {'param': 'value'}, sample_query)
|
||||
|
||||
assert result == 'plugin_result'
|
||||
mock_plugin_loader.invoke_tool.assert_called_once_with(
|
||||
'test_tool', {'param': 'value'}, sample_query
|
||||
)
|
||||
mock_plugin_loader.invoke_tool.assert_called_once_with('test_tool', {'param': 'value'}, sample_query)
|
||||
# MCP loader should not be called
|
||||
mock_mcp_loader.invoke_tool.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_calls_mcp_loader_when_plugin_not_found(
|
||||
self, mock_app_with_loaders, sample_query
|
||||
):
|
||||
async def test_execute_calls_mcp_loader_when_plugin_not_found(self, mock_app_with_loaders, sample_query):
|
||||
"""Test that execute_func_call uses MCP loader when plugin doesn't have tool."""
|
||||
toolmgr = get_toolmgr_module()
|
||||
|
||||
@@ -253,25 +218,16 @@ class TestToolManagerExecuteFuncCall:
|
||||
mock_mcp_loader.has_tool = AsyncMock(return_value=True)
|
||||
|
||||
manager = toolmgr.ToolManager(mock_app)
|
||||
manager.plugin_tool_loader = mock_plugin_loader
|
||||
manager.mcp_tool_loader = mock_mcp_loader
|
||||
self._wire_loaders(manager, mock_app, mock_plugin_loader, mock_mcp_loader)
|
||||
|
||||
result = await manager.execute_func_call(
|
||||
'test_tool',
|
||||
{'param': 'value'},
|
||||
sample_query
|
||||
)
|
||||
result = await manager.execute_func_call('test_tool', {'param': 'value'}, sample_query)
|
||||
|
||||
assert result == 'mcp_result'
|
||||
mock_mcp_loader.invoke_tool.assert_called_once_with(
|
||||
'test_tool', {'param': 'value'}, sample_query
|
||||
)
|
||||
mock_mcp_loader.invoke_tool.assert_called_once_with('test_tool', {'param': 'value'}, sample_query)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_raises_when_tool_not_found(
|
||||
self, mock_app_with_loaders, sample_query
|
||||
):
|
||||
"""Test that execute_func_call raises ValueError when tool not found."""
|
||||
async def test_execute_raises_when_tool_not_found(self, mock_app_with_loaders, sample_query):
|
||||
"""Test that execute_func_call raises ToolNotFoundError when tool not found."""
|
||||
toolmgr = get_toolmgr_module()
|
||||
|
||||
mock_app, mock_plugin_loader, mock_mcp_loader = mock_app_with_loaders
|
||||
@@ -279,20 +235,13 @@ class TestToolManagerExecuteFuncCall:
|
||||
mock_mcp_loader.has_tool = AsyncMock(return_value=False)
|
||||
|
||||
manager = toolmgr.ToolManager(mock_app)
|
||||
manager.plugin_tool_loader = mock_plugin_loader
|
||||
manager.mcp_tool_loader = mock_mcp_loader
|
||||
self._wire_loaders(manager, mock_app, mock_plugin_loader, mock_mcp_loader)
|
||||
|
||||
with pytest.raises(ValueError, match='未找到工具'):
|
||||
await manager.execute_func_call(
|
||||
'unknown_tool',
|
||||
{},
|
||||
sample_query
|
||||
)
|
||||
with pytest.raises(toolmgr.ToolNotFoundError, match='Tool not found: unknown_tool'):
|
||||
await manager.execute_func_call('unknown_tool', {}, sample_query)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_plugin_loader_checked_first(
|
||||
self, mock_app_with_loaders, sample_query
|
||||
):
|
||||
async def test_plugin_loader_checked_first(self, mock_app_with_loaders, sample_query):
|
||||
"""Test that plugin loader is checked before MCP loader."""
|
||||
toolmgr = get_toolmgr_module()
|
||||
|
||||
@@ -302,8 +251,7 @@ class TestToolManagerExecuteFuncCall:
|
||||
mock_mcp_loader.has_tool = AsyncMock(return_value=True)
|
||||
|
||||
manager = toolmgr.ToolManager(mock_app)
|
||||
manager.plugin_tool_loader = mock_plugin_loader
|
||||
manager.mcp_tool_loader = mock_mcp_loader
|
||||
self._wire_loaders(manager, mock_app, mock_plugin_loader, mock_mcp_loader)
|
||||
|
||||
await manager.execute_func_call('test_tool', {}, sample_query)
|
||||
|
||||
@@ -317,20 +265,30 @@ class TestToolManagerShutdown:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shutdown_calls_loader_shutdown(self):
|
||||
"""Test that shutdown calls shutdown on both loaders."""
|
||||
"""Test that shutdown calls shutdown on every registered loader."""
|
||||
toolmgr = get_toolmgr_module()
|
||||
|
||||
mock_app = Mock()
|
||||
mock_plugin_loader = Mock()
|
||||
mock_plugin_loader.shutdown = AsyncMock()
|
||||
mock_mcp_loader = Mock()
|
||||
mock_mcp_loader.shutdown = AsyncMock()
|
||||
|
||||
def _make_loader():
|
||||
loader = Mock()
|
||||
loader.shutdown = AsyncMock()
|
||||
return loader
|
||||
|
||||
mock_native_loader = _make_loader()
|
||||
mock_plugin_loader = _make_loader()
|
||||
mock_mcp_loader = _make_loader()
|
||||
mock_skill_loader = _make_loader()
|
||||
|
||||
manager = toolmgr.ToolManager(mock_app)
|
||||
manager.native_tool_loader = mock_native_loader
|
||||
manager.plugin_tool_loader = mock_plugin_loader
|
||||
manager.mcp_tool_loader = mock_mcp_loader
|
||||
manager.skill_tool_loader = mock_skill_loader
|
||||
|
||||
await manager.shutdown()
|
||||
|
||||
mock_native_loader.shutdown.assert_called_once()
|
||||
mock_plugin_loader.shutdown.assert_called_once()
|
||||
mock_mcp_loader.shutdown.assert_called_once()
|
||||
mock_mcp_loader.shutdown.assert_called_once()
|
||||
mock_skill_loader.shutdown.assert_called_once()
|
||||
|
||||
@@ -0,0 +1,499 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import os
|
||||
import tempfile
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
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.toolmgr import ToolManager
|
||||
|
||||
|
||||
class StubLoader:
|
||||
def __init__(
|
||||
self,
|
||||
tools: list[resource_tool.LLMTool] | None = None,
|
||||
invoke_result=None,
|
||||
catalog_source: str = 'mcp',
|
||||
catalog_source_name: str = 'fixture-server',
|
||||
):
|
||||
self._tools = tools or []
|
||||
self._invoke_result = invoke_result
|
||||
self._catalog_source = catalog_source
|
||||
self._catalog_source_name = catalog_source_name
|
||||
|
||||
async def get_tools(self, *_args, **_kwargs):
|
||||
return self._tools
|
||||
|
||||
async def get_tool_catalog(self, *_args, **_kwargs):
|
||||
return [
|
||||
{
|
||||
'name': tool.name,
|
||||
'description': tool.description,
|
||||
'human_desc': tool.human_desc,
|
||||
'parameters': tool.parameters,
|
||||
'source': self._catalog_source,
|
||||
'source_name': self._catalog_source_name,
|
||||
'source_id': self._catalog_source_name,
|
||||
}
|
||||
for tool in self._tools
|
||||
]
|
||||
|
||||
async def has_tool(self, name: str) -> bool:
|
||||
return any(tool.name == name for tool in self._tools)
|
||||
|
||||
async def invoke_tool(self, name: str, parameters: dict, query):
|
||||
return self._invoke_result(name, parameters, query) if callable(self._invoke_result) else self._invoke_result
|
||||
|
||||
async def shutdown(self):
|
||||
return None
|
||||
|
||||
|
||||
def make_tool(name: str) -> resource_tool.LLMTool:
|
||||
return resource_tool.LLMTool(
|
||||
name=name,
|
||||
human_desc=name,
|
||||
description=name,
|
||||
parameters={'type': 'object', 'properties': {}},
|
||||
func=lambda parameters: parameters,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_manager_omits_skill_authoring_tools_by_default():
|
||||
manager = ToolManager(SimpleNamespace())
|
||||
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()
|
||||
|
||||
assert [tool.name for tool in tools] == ['exec', 'plugin_tool', 'mcp_tool']
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_manager_includes_skill_authoring_tools_when_requested():
|
||||
manager = ToolManager(SimpleNamespace())
|
||||
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(include_skill_authoring=True)
|
||||
|
||||
assert [tool.name for tool in tools] == ['exec', 'activate', 'plugin_tool', 'mcp_tool']
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_manager_catalog_labels_tool_sources():
|
||||
manager = ToolManager(SimpleNamespace())
|
||||
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')],
|
||||
catalog_source='plugin',
|
||||
catalog_source_name='fixture-plugin',
|
||||
)
|
||||
manager.mcp_tool_loader = StubLoader([make_tool('mcp_tool')])
|
||||
|
||||
catalog = await manager.get_tool_catalog(include_skill_authoring=True)
|
||||
|
||||
assert [(item['name'], item['source'], item['source_name']) for item in catalog] == [
|
||||
('exec', 'builtin', 'LangBot'),
|
||||
('activate', 'skill', 'LangBot'),
|
||||
('plugin_tool', 'plugin', 'fixture-plugin'),
|
||||
('mcp_tool', 'mcp', 'fixture-server'),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_manager_routes_native_tool_calls():
|
||||
app = SimpleNamespace()
|
||||
manager = ToolManager(app)
|
||||
manager.native_tool_loader = StubLoader([make_tool('exec')], invoke_result={'backend': 'fake'})
|
||||
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')])
|
||||
|
||||
result = await manager.execute_func_call('exec', {'command': 'pwd'}, query=Mock())
|
||||
|
||||
assert result == {'backend': 'fake'}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_tool_loader_hides_tools_when_box_unavailable():
|
||||
loader = NativeToolLoader(SimpleNamespace(box_service=SimpleNamespace(available=False)))
|
||||
|
||||
assert await loader.get_tools() == []
|
||||
for tool_name in ('exec', 'read', 'write', 'edit', 'glob', 'grep'):
|
||||
assert await loader.has_tool(tool_name) is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
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}}),
|
||||
)
|
||||
loader = NativeToolLoader(SimpleNamespace(box_service=box_service, logger=Mock()))
|
||||
await loader.initialize()
|
||||
|
||||
tools = await loader.get_tools()
|
||||
|
||||
assert [tool.name for tool in tools] == ['exec', 'read', 'write', 'edit', 'glob', 'grep']
|
||||
for tool_name in ('exec', 'read', 'write', 'edit', 'glob', 'grep'):
|
||||
assert await loader.has_tool(tool_name) is True
|
||||
|
||||
|
||||
# ── 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)
|
||||
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
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_file():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
loader, _ = _make_loader_with_workspace(tmpdir)
|
||||
with open(os.path.join(tmpdir, 'hello.txt'), 'w') as f:
|
||||
f.write('hello world')
|
||||
|
||||
result = await loader.invoke_tool('read', {'path': '/workspace/hello.txt'}, _make_query())
|
||||
|
||||
assert result['ok'] is True
|
||||
assert result['content'] == 'hello world'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_nonexistent_file():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
loader, _ = _make_loader_with_workspace(tmpdir)
|
||||
|
||||
result = await loader.invoke_tool('read', {'path': '/workspace/no_such.txt'}, _make_query())
|
||||
|
||||
assert result['ok'] is False
|
||||
assert 'not found' in result['error'].lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_directory():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
loader, _ = _make_loader_with_workspace(tmpdir)
|
||||
os.makedirs(os.path.join(tmpdir, 'subdir'))
|
||||
with open(os.path.join(tmpdir, 'a.txt'), 'w') as f:
|
||||
f.write('a')
|
||||
|
||||
result = await loader.invoke_tool('read', {'path': '/workspace'}, _make_query())
|
||||
|
||||
assert result['ok'] is True
|
||||
assert result['is_directory'] is True
|
||||
assert 'a.txt' in result['content']
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_creates_file():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
loader, _ = _make_loader_with_workspace(tmpdir)
|
||||
|
||||
result = await loader.invoke_tool(
|
||||
'write', {'path': '/workspace/new.txt', 'content': 'new content'}, _make_query()
|
||||
)
|
||||
|
||||
assert result['ok'] is True
|
||||
with open(os.path.join(tmpdir, 'new.txt')) as f:
|
||||
assert f.read() == 'new content'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_creates_subdirectories():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
loader, _ = _make_loader_with_workspace(tmpdir)
|
||||
|
||||
result = await loader.invoke_tool(
|
||||
'write', {'path': '/workspace/sub/deep/file.txt', 'content': 'nested'}, _make_query()
|
||||
)
|
||||
|
||||
assert result['ok'] is True
|
||||
with open(os.path.join(tmpdir, 'sub', 'deep', 'file.txt')) as f:
|
||||
assert f.read() == 'nested'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_binary_file_as_base64_chunk():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
loader, _ = _make_loader_with_workspace(tmpdir)
|
||||
with open(os.path.join(tmpdir, 'blob.bin'), 'wb') as f:
|
||||
f.write(b'\x00\x01\x02\x03\x04')
|
||||
|
||||
result = await loader.invoke_tool(
|
||||
'read',
|
||||
{
|
||||
'path': '/workspace/blob.bin',
|
||||
'encoding': 'base64',
|
||||
'byte_offset': 1,
|
||||
'max_bytes': 2,
|
||||
},
|
||||
_make_query(),
|
||||
)
|
||||
|
||||
assert result['ok'] is True
|
||||
assert result['content'] == base64.b64encode(b'\x01\x02').decode('ascii')
|
||||
assert result['encoding'] == 'base64'
|
||||
assert result['byte_offset'] == 1
|
||||
assert result['length'] == 2
|
||||
assert result['size_bytes'] == 5
|
||||
assert result['has_more'] is True
|
||||
assert result['next_byte_offset'] == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_base64_file_append():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
loader, _ = _make_loader_with_workspace(tmpdir)
|
||||
|
||||
first = base64.b64encode(b'\x00\x01').decode('ascii')
|
||||
second = base64.b64encode(b'\x02\x03').decode('ascii')
|
||||
await loader.invoke_tool(
|
||||
'write',
|
||||
{'path': '/workspace/blob.bin', 'content': first, 'encoding': 'base64'},
|
||||
_make_query(),
|
||||
)
|
||||
result = await loader.invoke_tool(
|
||||
'write',
|
||||
{
|
||||
'path': '/workspace/blob.bin',
|
||||
'content': second,
|
||||
'encoding': 'base64',
|
||||
'mode': 'append',
|
||||
},
|
||||
_make_query(),
|
||||
)
|
||||
|
||||
assert result['ok'] is True
|
||||
with open(os.path.join(tmpdir, 'blob.bin'), 'rb') as f:
|
||||
assert f.read() == b'\x00\x01\x02\x03'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_base64_rejects_invalid_content():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
loader, _ = _make_loader_with_workspace(tmpdir)
|
||||
|
||||
result = await loader.invoke_tool(
|
||||
'write',
|
||||
{'path': '/workspace/blob.bin', 'content': 'not base64!', 'encoding': 'base64'},
|
||||
_make_query(),
|
||||
)
|
||||
|
||||
assert result['ok'] is False
|
||||
assert 'invalid base64' in result['error']
|
||||
assert not os.path.exists(os.path.join(tmpdir, 'blob.bin'))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_replaces_unique_string():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
loader, _ = _make_loader_with_workspace(tmpdir)
|
||||
with open(os.path.join(tmpdir, 'code.py'), 'w') as f:
|
||||
f.write('def foo():\n return 1\n')
|
||||
|
||||
result = await loader.invoke_tool(
|
||||
'edit',
|
||||
{'path': '/workspace/code.py', 'old_string': 'return 1', 'new_string': 'return 42'},
|
||||
_make_query(),
|
||||
)
|
||||
|
||||
assert result['ok'] is True
|
||||
with open(os.path.join(tmpdir, 'code.py')) as f:
|
||||
assert f.read() == 'def foo():\n return 42\n'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_rejects_ambiguous_match():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
loader, _ = _make_loader_with_workspace(tmpdir)
|
||||
with open(os.path.join(tmpdir, 'dup.txt'), 'w') as f:
|
||||
f.write('aaa\naaa\n')
|
||||
|
||||
result = await loader.invoke_tool(
|
||||
'edit',
|
||||
{'path': '/workspace/dup.txt', 'old_string': 'aaa', 'new_string': 'bbb'},
|
||||
_make_query(),
|
||||
)
|
||||
|
||||
assert result['ok'] is False
|
||||
assert '2' in result['error']
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_rejects_missing_string():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
loader, _ = _make_loader_with_workspace(tmpdir)
|
||||
with open(os.path.join(tmpdir, 'x.txt'), 'w') as f:
|
||||
f.write('hello')
|
||||
|
||||
result = await loader.invoke_tool(
|
||||
'edit',
|
||||
{'path': '/workspace/x.txt', 'old_string': 'nope', 'new_string': 'yes'},
|
||||
_make_query(),
|
||||
)
|
||||
|
||||
assert result['ok'] is False
|
||||
assert 'not found' in result['error'].lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_path_escape_blocked():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
loader, _ = _make_loader_with_workspace(tmpdir)
|
||||
|
||||
with pytest.raises(ValueError, match='escapes'):
|
||||
await loader.invoke_tool('read', {'path': '/workspace/../../etc/passwd'}, _make_query())
|
||||
|
||||
|
||||
@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
|
||||
|
||||
assert await is_box_backend_available(SimpleNamespace()) is False
|
||||
assert await is_box_backend_available(SimpleNamespace(box_service=SimpleNamespace(available=False))) is False
|
||||
|
||||
unavailable_backend = SimpleNamespace(
|
||||
available=True,
|
||||
get_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')),
|
||||
)
|
||||
assert await is_box_backend_available(SimpleNamespace(box_service=failing_backend)) is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_file_supports_offset_limit_and_truncation_metadata():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
loader, _ = _make_loader_with_workspace(tmpdir)
|
||||
with open(os.path.join(tmpdir, 'lines.txt'), 'w', encoding='utf-8') as f:
|
||||
f.write('one\ntwo\nthree\nfour\n')
|
||||
|
||||
result = await loader.invoke_tool(
|
||||
'read',
|
||||
{'path': '/workspace/lines.txt', 'offset': 2, 'limit': 2},
|
||||
_make_query(),
|
||||
)
|
||||
|
||||
assert result == {
|
||||
'ok': True,
|
||||
'content': 'two\nthree',
|
||||
'truncated': True,
|
||||
'truncated_by': 'lines',
|
||||
'start_line': 2,
|
||||
'end_line': 3,
|
||||
'next_offset': 4,
|
||||
'max_lines': 2,
|
||||
'max_bytes': 50 * 1024,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_file_handles_line_larger_than_byte_limit():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
loader, _ = _make_loader_with_workspace(tmpdir)
|
||||
with open(os.path.join(tmpdir, 'long-line.txt'), 'w', encoding='utf-8') as f:
|
||||
f.write('abcdef\n')
|
||||
|
||||
result = await loader.invoke_tool(
|
||||
'read',
|
||||
{'path': '/workspace/long-line.txt', 'max_bytes': 3},
|
||||
_make_query(),
|
||||
)
|
||||
|
||||
assert result['ok'] is True
|
||||
assert result['truncated'] is True
|
||||
assert result['truncated_by'] == 'bytes'
|
||||
assert result['next_offset'] == 1
|
||||
assert 'exceeds the 3B read limit' in result['content']
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_exec_result_is_capped_and_exposes_preview_metadata():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
box_service = SimpleNamespace(
|
||||
available=True,
|
||||
default_workspace=tmpdir,
|
||||
execute_tool=AsyncMock(
|
||||
return_value={
|
||||
'ok': True,
|
||||
'stdout': 'a' * 60000,
|
||||
'stderr': 'b' * 60000,
|
||||
'exit_code': 0,
|
||||
}
|
||||
),
|
||||
)
|
||||
loader = NativeToolLoader(SimpleNamespace(box_service=box_service, logger=Mock()))
|
||||
|
||||
result = await loader.invoke_tool('exec', {'command': 'python -V'}, _make_query())
|
||||
|
||||
assert result['ok'] is True
|
||||
assert len(result['stdout'].encode('utf-8')) == 50 * 1024
|
||||
assert len(result['stderr'].encode('utf-8')) == 50 * 1024
|
||||
assert len(result['preview'].encode('utf-8')) == 50 * 1024
|
||||
assert result['stdout_truncated'] is True
|
||||
assert result['stderr_truncated'] is True
|
||||
assert result['truncated'] is True
|
||||
assert result['truncated_by'] == 'bytes'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_glob_caps_match_count_and_returns_preview():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
loader, _ = _make_loader_with_workspace(tmpdir)
|
||||
for index in range(105):
|
||||
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'] == 105
|
||||
assert len(result['matches']) == 100
|
||||
assert result['preview'] == '\n'.join(result['matches'])
|
||||
assert result['truncated'] is True
|
||||
assert result['truncated_by'] == 'matches'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_grep_reports_invalid_regex_and_truncates_long_matching_lines():
|
||||
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('needle ' + ('x' * 600) + '\n')
|
||||
|
||||
invalid = await loader.invoke_tool('grep', {'path': '/workspace', 'pattern': '['}, _make_query())
|
||||
result = await loader.invoke_tool('grep', {'path': '/workspace', 'pattern': 'needle'}, _make_query())
|
||||
|
||||
assert invalid['ok'] is False
|
||||
assert 'Invalid regex' in invalid['error']
|
||||
assert result['ok'] is True
|
||||
assert result['truncated'] is True
|
||||
assert result['truncated_by'] == 'line'
|
||||
assert result['matches'][0]['file'] == '/workspace/data.txt'
|
||||
assert result['matches'][0]['content'].endswith('... [truncated]')
|
||||
Reference in New Issue
Block a user